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..3eb54cad62f --- /dev/null +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -0,0 +1,74 @@ +--- +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 queued response origins before restoring shared or + executor state, including embedded child checkpoints. Submitted and already + queued responses share the same batch validation before pending requests are consumed. +- 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. + +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 ad8a87b1080..b95cf65dcc1 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -333,6 +333,9 @@ 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 checkpoint restoration** - Stored and embedded restores preflight a detached checkpoint tree before + replacing live state or queues. `WorkflowExecutor` delegates its private validation hook to its child's runner, + so each child uses its own response-origin policy and migration metadata stays in the candidate snapshot. - **`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 cee5d96abd7..5181a867abb 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -332,6 +332,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, @@ -418,6 +419,7 @@ async def _run_stream_impl( all_updates: list[AgentResponseUpdate] = [] 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, @@ -439,6 +441,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, @@ -453,6 +456,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. @@ -470,19 +474,23 @@ 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, - ): + ) + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] + 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, ) + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] + _ = await restore_result if not input_messages: logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") return @@ -500,43 +508,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..7f096c57427 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -17,6 +17,12 @@ # 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" + # 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/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 7c1a6ffd524..6a9947a72ba 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -7,7 +7,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 @@ -31,6 +31,8 @@ logger = logging.getLogger(__name__) +_ResponseBatchValidation = tuple[Callable[[Sequence[WorkflowEvent]], None], Sequence[WorkflowEvent]] + # region Executor class Executor(RequestInfoMixin, DictConvertible): @@ -533,6 +535,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 [] + + 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/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index f41183cb1c2..7703c7f56f3 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -1,10 +1,11 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import copy import logging import warnings from collections import defaultdict -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, Callable, Sequence from typing import TYPE_CHECKING, Any, cast from ..exceptions import ( @@ -20,6 +21,7 @@ from ._runner_context import ( RunnerContext, WorkflowMessage, + _get_queued_response_info_events, # pyright: ignore[reportPrivateUsage] ) from ._state import State @@ -83,6 +85,7 @@ def __init__( # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None + self._response_batch_validator: Callable[[Sequence[WorkflowEvent]], None] | None = None @property def context(self) -> RunnerContext: @@ -319,12 +322,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) + 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 +381,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) + 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 +393,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 + def _validate_checkpoint_tree( + self, checkpoint: WorkflowCheckpoint, delegated_requests: Sequence[WorkflowEvent] = () + ) -> None: + from ._workflow_executor import SubWorkflowResponseMessage + + 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) + 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) + # Wrapped replies retain child-local source IDs, so use the parent graph to route them. + for batch in checkpoint.messages.values(): + for message in batch: + if not isinstance(message.data, SubWorkflowResponseMessage): + continue + targets = { + edge.target_id + for runner in self._edge_runner_map.get(message.source_id, []) + for edge in runner._edge_group.edges # pyright: ignore[reportPrivateUsage] + if edge.source_id == message.source_id + and (message.target_id is None or edge.target_id == message.target_id) + } + for target_id in targets: + delegated_by_executor[target_id].append(message.data.source_event) + for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): + 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 +461,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 +486,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..82b91558e73 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,15 @@ 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 RunnerContext(Protocol): """Protocol for the execution context used by the runner. @@ -128,6 +137,10 @@ async def has_messages(self) -> bool: """ ... + async def get_queued_response_info_events(self) -> list[WorkflowEvent[Any]]: + """Get original requests for queued responses without consuming messages.""" + ... + async def add_event(self, event: WorkflowEvent) -> None: """Add an event to the execution context. @@ -611,6 +624,10 @@ 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 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/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index c60617dabd5..85777174274 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -13,6 +13,7 @@ import weakref from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Collection, Mapping, Sequence from contextlib import aclosing +from copy import deepcopy from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload @@ -35,6 +36,8 @@ 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_RUN_KWARGS_KEY, ) from ._edge import ( @@ -48,7 +51,10 @@ WorkflowRunState, _framework_event_origin, # type: ignore ) -from ._executor import Executor +from ._executor import ( + Executor, + _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] +) from ._model_utils import DictConvertible from ._runner import RunnerImpl from ._runner_context import RunnerContext, WorkflowMessage @@ -411,6 +417,9 @@ 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: Callable[[Sequence[WorkflowEvent]], None] | None = None + self._runner._response_batch_validator = self._validate_response_batch # pyright: ignore[reportPrivateUsage] + 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 +438,10 @@ 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 to_dict(self) -> dict[str, Any]: """Serialize the workflow definition into a JSON-ready dictionary.""" data: dict[str, Any] = { @@ -761,6 +774,9 @@ async def _execute_with_message_or_checkpoint( f"Expected input types: {start_executor.input_types}." ) + 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 +887,10 @@ 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) + # 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 @@ -1130,6 +1150,35 @@ 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: + self._response_batch_validator(requests) + + async def _validate_response_requests(self, requests: Sequence[WorkflowEvent]) -> None: + validations = await self._get_response_validation_batches(requests) + # Preserve shared request data while preflighting the whole tree before live migrations. + detached_batches = deepcopy([batch for _, batch in validations]) + for (validator, _), batch in zip(validations, detached_batches, strict=True): + validator(batch) + for validator, batch in validations: + validator(batch) + + async def _get_response_validation_batches( + self, requests: Sequence[WorkflowEvent] + ) -> list[_ResponseBatchValidation]: + queued_requests = await self._runner.context.get_queued_response_info_events() + batch = [*queued_requests, *requests] + validations: list[_ResponseBatchValidation] = [] + if self._response_batch_validator is not None: + validations.append((self._validate_response_batch, batch)) + for executor in self.executors.values(): + delegated = [request for request in batch if request.source_executor_id == 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 +1197,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: diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 7313e91f322..780b5b74554 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -3,7 +3,7 @@ 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 +24,11 @@ WorkflowRunState, _framework_event_origin, # type: ignore[reportPrivateUsage] ) -from ._executor import Executor, handler +from ._executor import ( + Executor, + _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] + handler, +) from ._request_info_mixin import response_handler from ._runner_context import WorkflowMessage from ._typing_utils import is_instance_of @@ -590,6 +594,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 + 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: + 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.""" diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 6e13a414f8b..642dde75596 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -654,8 +654,12 @@ 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"] + assert [[event.request_id for event in batch] for batch in validated_requests] == [[request_id]] + assert validated_requests[0][0].data == "approve" async def test_memory_checkpoint_storage_roundtrip_json_native_types(): diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 97ee0433694..22b60ed2f72 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, @@ -29,11 +31,13 @@ WorkflowBuilder, 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 +198,114 @@ 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("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..d6c26bbee8c 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -24,6 +24,17 @@ series because response caching and default-header integration use its internals ## MCP Handler Lifetimes `DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider`. +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. +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. +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 +43,12 @@ 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. +Without a provider, `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. + 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 +60,11 @@ 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. 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 +79,30 @@ 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 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 already queued responses +before any pending request is consumed. +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..428e92755eb 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -91,3 +91,46 @@ This intentionally incurs connection and initialization overhead and does not retain server session state between provider-backed invocations. 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..9ed8200fd12 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 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,35 @@ } +def _validate_response_batch( + requests: Sequence[WorkflowEvent[Any]], *, migrate_legacy_request_scopes: bool = False +) -> 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.") + for data, scope in scoped_requests: + data.workflow_session_id = scope + + +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 +182,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 +205,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 +222,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 +285,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 1d642c91c3a..7d56f8f874d 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 @@ -488,6 +495,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) @@ -524,6 +532,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..1784a409ce4 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,21 @@ 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 ._errors import DeclarativeWorkflowError + if TYPE_CHECKING: from agent_framework import Content + from agent_framework._workflows._state import State __all__ = [ "ClientProvider", @@ -53,6 +58,74 @@ 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: + agent_session_id = state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) + if agent_session_id is not 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: + if state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) is not None: + get_or_create_workflow_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) + + +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) + + +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 +147,9 @@ 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. """ server_url: str @@ -82,6 +158,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,6 +231,23 @@ 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: @@ -161,12 +255,16 @@ class DefaultMCPToolHandler: Without a ``client_provider``, 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 @@ -198,10 +296,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 MCP clients in no-provider mode, + 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 invocations. """ LIST_TOOLS_TOOL_NAME: ClassVar[str] = "tools/list" @@ -229,14 +329,17 @@ 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]] = {} + 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: provider-backed calls never share entries. self._active_invocations: set[asyncio.Future[None]] = set() # Keep ancestry so a completed nested call cannot hide an active parent @@ -247,6 +350,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. @@ -304,7 +408,9 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: return await self._invoke_entry(entry, invocation) finally: - if completion is not None: + if self._client_provider is None 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) @@ -442,7 +548,7 @@ async def aclose(self) -> None: 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 @@ -458,17 +564,53 @@ async def aclose(self) -> None: "DefaultMCPToolHandler.aclose() cannot be called from an active provider-backed 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 +622,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 +653,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 +663,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 +897,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 +906,15 @@ 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 ( + 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 +927,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 +942,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,16 +971,18 @@ 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 @@ -683,4 +994,10 @@ def _cache_key( 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 or "", + server_url, + server_label or "", + connection_name or "", + headers_hash, + ) 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..5ae6a13be37 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( @@ -830,6 +834,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 +960,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 +1108,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 +1513,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 +1624,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 +1761,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 +2073,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 3982831a9fc..211a3b698d2 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..6e8f81775a2 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,9 @@ """ import sys +from pathlib import Path from typing import Any +from unittest.mock import AsyncMock, MagicMock import httpx2 as httpx import pytest @@ -26,7 +28,7 @@ reason="PowerFx engine not available (requires dotnet runtime)", ) -from agent_framework import Content, Message # noqa: E402 +from agent_framework import AgentSession, Content, Message # noqa: E402 from agent_framework.exceptions import ToolExecutionException # noqa: E402 from agent_framework_declarative._workflows import ( # noqa: E402 @@ -164,6 +166,668 @@ 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("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("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 +898,1020 @@ 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_preflight_preserves_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, + 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 + from agent_framework.exceptions import WorkflowCheckpointException + + 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 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 + assert recorder.value == "live" + assert recorder.restore_calls == 0 + assert all(workflow._runner.state.get("uncommitted") == "keep" for workflow in workflows) + else: + await restore() + _ = [event async for event in parent._runner.run_until_convergence()] + assert all(workflow._runner.state.get("restore_marker") == "snapshot" for workflow in workflows) + assert recorder.value == "snapshot" + assert recorder.restore_calls == 1 + 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 encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before + assert mcp_handler.call_count == (2 if mode in {"migrate", "valid"} else 0) + + @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, + 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): + queued = await leaf._runner_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 ------------------------------------------------ @@ -445,6 +2123,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 +2134,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 +2153,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 +2166,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 +2267,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,