Compare commits

...

23 Commits

Author SHA1 Message Date
Tao Chen f3f1d4f969 fix unit tests 2026-06-23 16:28:16 -07:00
Tao Chen fee3840579 Merge branch 'main' into local-branch-python-add-reset-to-workflow 2026-06-23 11:39:20 -07:00
Tao Chen 1e7b9f60b0 Add runner back to public API for backward comp but with deprecation warning 2026-06-17 14:07:27 -07:00
Tao Chen 0aa80f686f Capture initial checkpoint in memory 2026-06-17 13:48:13 -07:00
Tao Chen d4cf991fea Remove merge error 2026-06-16 11:52:53 -07:00
Tao Chen c7aac7c8ff Merge branch 'main' into local-branch-python-add-reset-to-workflow 2026-06-16 11:50:16 -07:00
Tao Chen ea052ab511 Merge branch 'main' into local-branch-python-add-reset-to-workflow 2026-06-15 09:50:34 -07:00
Tao Chen d08cec5379 Addres race condition when stream is dropped midway 2026-06-15 09:40:21 -07:00
Tao Chen 422f7e7382 Address comments 2026-06-11 16:54:42 -07:00
Tao Chen 96af0cd15a Merge branch 'main' into local-branch-python-add-reset-to-workflow 2026-06-11 15:28:27 -07:00
Tao Chen b0d0224ed4 Address comments 2026-06-11 15:27:36 -07:00
Tao Chen ed27241543 Add tests 2026-06-11 11:43:41 -07:00
Tao Chen 9da83347c8 Remove reset 2026-06-11 11:14:36 -07:00
Tao Chen 6534a739d0 Add create checkpoint to workflow interface 2026-06-10 22:20:10 -07:00
Tao Chen 0e3831192a Fix checkpoint ancestry bug 2026-06-10 16:36:15 -07:00
Tao Chen 47012f1dcf Address comments 2026-06-09 11:02:17 -07:00
Tao Chen 4e623d561f Merge branch 'main' into local-branch-python-add-reset-to-workflow 2026-06-09 09:38:45 -07:00
Tao Chen 6eb85477fc Remove lifecycle flag 2026-06-09 09:27:21 -07:00
Tao Chen 568afdd293 Fix tests and address comments 2026-06-08 16:30:59 -07:00
Tao Chen 598ad231ba Add sample 2026-06-08 16:07:26 -07:00
Tao Chen 5910a6f869 Add reset to hosted workflow 2026-06-08 14:57:01 -07:00
Tao Chen 65522bdbee Add reset to workflow 2026-06-08 13:42:42 -07:00
Tao Chen c5e6a7797f Move runner state management out of Workflow 2026-06-05 16:29:19 -07:00
15 changed files with 1107 additions and 264 deletions
@@ -9,7 +9,7 @@ integrations, many of which are lazy-loaded from optional packages.
"""
import importlib.metadata
from typing import Final
from typing import TYPE_CHECKING, Any, Final
try:
_version = importlib.metadata.version(__name__)
@@ -264,6 +264,7 @@ from ._workflows._agent_executor import (
)
from ._workflows._agent_utils import resolve_agent_id
from ._workflows._checkpoint import (
CheckpointID,
CheckpointStorage,
FileCheckpointStorage,
InMemoryCheckpointStorage,
@@ -307,7 +308,6 @@ from ._workflows._functional import (
workflow,
)
from ._workflows._request_info_mixin import response_handler
from ._workflows._runner import Runner
from ._workflows._runner_context import (
InProcRunnerContext,
RunnerContext,
@@ -405,6 +405,7 @@ __all__ = [
"ChatResponse",
"ChatResponseUpdate",
"CheckResult",
"CheckpointID",
"CheckpointStorage",
"ClassSkill",
"CompactionProvider",
@@ -618,3 +619,20 @@ __all__ = [
"validate_workflow_graph",
"workflow",
]
if TYPE_CHECKING:
from ._workflows._runner import Runner
def __getattr__(name: str) -> Any:
"""Lazily resolve deprecated public names, emitting a ``DeprecationWarning``.
``Runner`` remains importable from ``agent_framework`` for backward
compatibility but is deprecated and slated for removal from the public API.
"""
if name == "Runner":
from ._workflows._runner import Runner, warn_runner_deprecated
warn_runner_deprecated()
return Runner
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -3,6 +3,7 @@
import asyncio
import contextlib
import logging
import warnings
from collections import defaultdict
from collections.abc import AsyncGenerator, Sequence
from typing import Any
@@ -10,7 +11,7 @@ from typing import Any
from ..exceptions import (
WorkflowCheckpointException,
WorkflowConvergenceException,
WorkflowRunnerException,
WorkflowException,
)
from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint
from ._const import EXECUTOR_STATE_KEY
@@ -27,6 +28,21 @@ from ._state import State
logger = logging.getLogger(__name__)
def warn_runner_deprecated() -> None:
"""Emit a deprecation warning when ``Runner`` is accessed from the public API.
``Runner`` remains importable from ``agent_framework`` for backward
compatibility, but it is intended for internal use only and will be removed
from the public API in a future version.
"""
warnings.warn(
"`Runner` is deprecated and will be removed from the public API in a future version. "
"It is intended for internal use only.",
DeprecationWarning,
stacklevel=3,
)
class Runner:
"""A class to run a workflow in Pregel supersteps."""
@@ -63,99 +79,123 @@ class Runner:
self._iteration = 0
self._max_iterations = max_iterations
self._state = state
self._running = False
self._resumed_from_checkpoint = False # Track whether we resumed
# Checkpointing related attributes
self._resumed_from_checkpoint = False
self._previous_checkpoint_id: CheckpointID | None = None
@property
def context(self) -> RunnerContext:
"""Get the workflow context."""
"""Get the runner context for message, event, and checkpoint handling."""
return self._ctx
@property
def state(self) -> State:
"""Get the shared state for the workflow."""
return self._state
def reset_iteration_count(self) -> None:
"""Reset the iteration count to zero."""
"""Reset the iteration count to zero.
This is useful when the workflow resumes from a new set of messages.
Note:
When a workflow is resumed from a response (for a request_info_event)
or a checkpoint, the iteration count is normally NOT reset.
"""
self._iteration = 0
def reset_runtime_state(
self,
*,
iteration: int = 0,
previous_checkpoint_id: CheckpointID | None = None,
resumed_from_checkpoint: bool = False,
) -> None:
"""Reset runner runtime bookkeeping to a known baseline.
Args:
iteration: Iteration value to restore.
previous_checkpoint_id: Checkpoint parent pointer for subsequent saves.
resumed_from_checkpoint: Whether to treat next run as resumed.
"""
self._iteration = iteration
self._previous_checkpoint_id = previous_checkpoint_id
self._resumed_from_checkpoint = resumed_from_checkpoint
async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
"""Run the workflow until no more messages are sent."""
if self._running:
raise WorkflowRunnerException("Runner is already running.")
# Emit any events already produced prior to entering loop
if await self._ctx.has_events():
logger.info("Yielding pre-loop events")
for event in await self._ctx.drain_events():
yield event
self._running = True
previous_checkpoint_id: CheckpointID | None = None
try:
# Emit any events already produced prior to entering loop
if await self._ctx.has_events():
logger.info("Yielding pre-loop events")
for event in await self._ctx.drain_events():
yield event
# Create a checkpoint before a run starts. Checkpoints are usually considered to be created at the
# end of an iteration, we can think of this checkpoint as being created at the end of "superstep 0"
# which captures the states after which the start executor has run. Note that we execute the start
# executor outside of the main iteration loop.
if await self._ctx.has_messages() and not self._resumed_from_checkpoint:
await self.create_checkpoint_if_enabled()
# Create the first checkpoint. Checkpoints are usually considered to be created at the end of an iteration,
# we can think of the first checkpoint as being created at the end of a "superstep 0" which captures the
# states after which the start executor has run. Note that we execute the start executor outside of the
# main iteration loop.
if await self._ctx.has_messages() and not self._resumed_from_checkpoint:
previous_checkpoint_id = await self._create_checkpoint_if_enabled(previous_checkpoint_id)
while self._iteration < self._max_iterations:
logger.info(f"Starting superstep {self._iteration + 1}")
yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
while self._iteration < self._max_iterations:
logger.info(f"Starting superstep {self._iteration + 1}")
yield WorkflowEvent.superstep_started(iteration=self._iteration + 1)
# Run iteration concurrently with live event streaming: we poll
# for new events while the iteration coroutine progresses.
iteration_task = asyncio.create_task(self._run_iteration())
try:
while not iteration_task.done():
try:
# Wait briefly for any new event; timeout allows progress checks
event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
yield event
except asyncio.TimeoutError:
# Periodically continue to let iteration advance
continue
except asyncio.CancelledError:
# Propagate cancellation to the iteration task to avoid orphaned work
iteration_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await iteration_task
raise
# Propagate errors from iteration, but first surface any pending events
try:
# Run iteration concurrently with live event streaming: we poll
# for new events while the iteration coroutine progresses.
iteration_task = asyncio.create_task(self._run_iteration())
try:
while not iteration_task.done():
try:
# Wait briefly for any new event; timeout allows progress checks
event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
yield event
except asyncio.TimeoutError:
# Periodically continue to let iteration advance
continue
except asyncio.CancelledError:
# Propagate cancellation to the iteration task to avoid orphaned work
iteration_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await iteration_task
except Exception:
# Make sure failure-related events (like ExecutorFailedEvent) are surfaced
if await self._ctx.has_events():
for event in await self._ctx.drain_events():
yield event
raise
self._iteration += 1
raise
# Drain any straggler events emitted at tail end
# Propagate errors from iteration, but first surface any pending events
try:
await iteration_task
except Exception:
# Make sure failure-related events (like ExecutorFailedEvent) are surfaced
if await self._ctx.has_events():
for event in await self._ctx.drain_events():
yield event
raise
self._iteration += 1
logger.info(f"Completed superstep {self._iteration}")
# Drain any straggler events emitted at tail end
if await self._ctx.has_events():
for event in await self._ctx.drain_events():
yield event
# Commit pending state changes at superstep boundary
self._state.commit()
logger.info(f"Completed superstep {self._iteration}")
# Create checkpoint after each superstep iteration
previous_checkpoint_id = await self._create_checkpoint_if_enabled(previous_checkpoint_id)
# Commit pending state changes at superstep boundary
self._state.commit()
yield WorkflowEvent.superstep_completed(iteration=self._iteration)
# Create checkpoint after each superstep iteration
await self.create_checkpoint_if_enabled()
# Check for convergence: no more messages to process
if not await self._ctx.has_messages():
break
yield WorkflowEvent.superstep_completed(iteration=self._iteration)
if self._iteration >= self._max_iterations and await self._ctx.has_messages():
raise WorkflowConvergenceException(f"Runner did not converge after {self._max_iterations} iterations.")
# Check for convergence: no more messages to process
if not await self._ctx.has_messages():
break
logger.info(f"Workflow completed after {self._iteration} supersteps")
self._resumed_from_checkpoint = False # Reset resume flag for next run
finally:
self._running = False
logger.info(f"Workflow completed after {self._iteration} supersteps")
self._resumed_from_checkpoint = False # Reset resume flag for next run
if self._iteration >= self._max_iterations and await self._ctx.has_messages():
raise WorkflowConvergenceException(f"Runner did not converge after {self._max_iterations} iterations.")
async def _run_iteration(self) -> None:
"""Run a single iteration of the workflow.
@@ -209,40 +249,121 @@ class Runner:
]
await asyncio.gather(*tasks)
async def _create_checkpoint_if_enabled(self, previous_checkpoint_id: CheckpointID | None) -> CheckpointID | None:
async def _prepare_checkpoint_state(self) -> None:
"""Persist executor snapshots into committed shared state.
This is used by checkpoint capture paths that need a complete, restorable
state payload without necessarily writing to a checkpoint storage backend.
"""
await self._save_executor_states()
self._state.commit()
async def capture_checkpoint_object(self, *, metadata: dict[str, Any] | None = None) -> WorkflowCheckpoint:
"""Capture the current runner state as an in-memory checkpoint object.
Persists executor snapshots into committed state and builds a
``WorkflowCheckpoint`` from the current committed state. The checkpoint is
not written to any storage backend; the caller owns its lifetime (for
example, the workflow's captured initial checkpoint used by reset).
This is only valid when the runner is quiescent: it rejects capture when
in-flight executor messages or pending request_info events are present,
since those represent mid-run state that would not form a clean baseline.
Args:
metadata: Optional metadata to attach to the checkpoint.
Returns:
A ``WorkflowCheckpoint`` snapshot of the current runner state.
Raises:
WorkflowException: If in-flight messages or pending requests are present.
"""
if await self._ctx.has_messages():
raise WorkflowException("Cannot capture checkpoint while in-flight messages are present.")
pending_requests = await self._ctx.get_pending_request_info_events()
if pending_requests:
raise WorkflowException("Cannot capture checkpoint while pending requests are present.")
await self._prepare_checkpoint_state()
return WorkflowCheckpoint(
workflow_name=self._workflow_name,
graph_signature_hash=self._graph_signature_hash,
previous_checkpoint_id=None,
messages={},
state=self._state.export_state(),
pending_request_info_events={},
iteration_count=0,
metadata=metadata or {},
)
async def create_checkpoint_if_enabled(self) -> None:
"""Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
if not self._ctx.has_checkpointing():
return None
return
try:
# Save executor states into the shared state before creating the checkpoint,
# so that they are included in the checkpoint payload.
await self._save_executor_states()
# `on_checkpoint_save()` writes via State.set(), which stages values in the
# pending buffer. Checkpoints serialize committed state only, so commit here
# to ensure executor snapshots are captured in this checkpoint.
self._state.commit()
# Save executor states into committed state before creating the checkpoint.
await self._prepare_checkpoint_state()
checkpoint_id = await self._ctx.create_checkpoint(
self._workflow_name,
self._graph_signature_hash,
self._state,
previous_checkpoint_id,
self._previous_checkpoint_id,
self._iteration,
)
logger.info(f"Created checkpoint: {checkpoint_id}")
return checkpoint_id
logger.info(
"Created checkpoint: %s with parent checkpoint at iteration %d: %s",
checkpoint_id,
self._iteration,
self._previous_checkpoint_id,
)
self._previous_checkpoint_id = checkpoint_id
except Exception as e:
logger.warning(f"Failed to create checkpoint: {e}")
return None
logger.warning(
"Failed to create checkpoint at iteration %d: %s. "
"Note that this does not fail the workflow run. "
"The next successfully-created checkpoint will be parented to the last successful checkpoint: %s",
self._iteration,
e,
self._previous_checkpoint_id,
)
async def restore_from_checkpoint_object(self, checkpoint: WorkflowCheckpoint) -> None:
"""Restore runner state from an in-memory checkpoint object.
Unlike :meth:`restore_from_checkpoint`, this does not load from storage or
validate the graph signature; it applies a checkpoint that the caller already
holds (for example, the workflow's captured initial checkpoint used by reset).
This clears any runtime checkpoint storage override and resets the context for a
fresh run, then restores shared state, executor snapshots, and runtime bookkeeping
from the checkpoint.
Args:
checkpoint: The checkpoint whose state should be restored.
"""
self._ctx.clear_runtime_checkpoint_storage()
self._ctx.reset_for_new_run()
self._state.clear()
self._state.import_state(checkpoint.state)
await self._restore_executor_states()
self.reset_runtime_state(
iteration=checkpoint.iteration_count,
previous_checkpoint_id=checkpoint.previous_checkpoint_id,
resumed_from_checkpoint=False,
)
async def restore_from_checkpoint(
self,
checkpoint_id: CheckpointID,
checkpoint_storage: CheckpointStorage | None = None,
) -> None:
"""Restore workflow state from a checkpoint.
"""Restore the runner from a checkpoint.
Args:
checkpoint_id: The ID of the checkpoint to restore from
@@ -290,7 +411,7 @@ class Runner:
# Apply the checkpoint to the context
await self._ctx.apply_checkpoint(checkpoint)
# Mark the runner as resumed
self._mark_resumed(checkpoint.iteration_count)
self._mark_resumed(checkpoint)
logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
except WorkflowCheckpointException:
@@ -356,13 +477,14 @@ class Runner:
return parsed
def _mark_resumed(self, iteration: int) -> None:
def _mark_resumed(self, checkpoint: WorkflowCheckpoint) -> None:
"""Mark the runner as having resumed from a checkpoint.
Optionally set the current iteration and max iterations.
"""
self._resumed_from_checkpoint = True
self._iteration = iteration
self._iteration = checkpoint.iteration_count
self._previous_checkpoint_id = checkpoint.checkpoint_id
async def _set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:
"""Store executor state in state under a reserved key.
@@ -403,12 +403,14 @@ class InProcRunnerContext:
def reset_for_new_run(self) -> None:
"""Reset the context for a new workflow run.
This clears messages, events, and resets streaming flag.
Runtime checkpoint storage is NOT cleared here as it's managed at the workflow level.
Clears messages, the pending event queue, the pending request_info
correlation map, and the streaming flag. Runtime checkpoint storage is
NOT cleared here as it's managed at the workflow level.
"""
self._messages.clear()
# Clear any pending events (best-effort) by recreating the queue
self._event_queue = asyncio.Queue()
self._pending_request_info_events.clear()
self._streaming = False # Reset streaming flag
async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
@@ -11,14 +11,16 @@ import logging
import types
import uuid
import warnings
import weakref
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal, overload
from .._sessions import ContextProvider
from .._types import ResponseStream
from ..exceptions import WorkflowException
from ..observability import OtelAttr, capture_exception, create_workflow_span
from ._checkpoint import CheckpointStorage
from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
from ._edge import (
EdgeGroup,
@@ -346,25 +348,33 @@ class Workflow(DictConvertible):
# Store non-serializable runtime objects as private attributes
self._runner_context = runner_context
self._runner_context.set_yield_output_classifier(self._output_designation.classify)
self._state = State()
self._runner: Runner = Runner(
self.edge_groups,
self.executors,
self._state,
State(),
runner_context,
self.name,
self.graph_signature_hash,
max_iterations=max_iterations,
)
# Flag to prevent concurrent workflow executions
self._is_running = False
# Current run-level status of this workflow instance. Updated in lockstep with
# the status events emitted from `_run_workflow_with_tracing`. Defaults to IDLE
# for a freshly built workflow that has not yet been run.
self._status: WorkflowRunState = WorkflowRunState.IDLE
# Weak reference to the in-flight run's ``ResponseStream``. Used as the single
# concurrency lock: if the previous stream is still alive, ``run()`` rejects a
# new run synchronously (before any await). When the stream is fully consumed
# ``_run_core``'s finally clears this; if the caller drops the stream without
# 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
# In-memory initial checkpoint captured from the just-built workflow state.
# This is internal-only and used by ``reset()``.
self._initial_checkpoint: WorkflowCheckpoint | None = None
@property
def status(self) -> WorkflowRunState:
"""Return the current run-level status of this workflow instance.
@@ -376,16 +386,6 @@ class Workflow(DictConvertible):
"""
return self._status
def _ensure_not_running(self) -> None:
"""Ensure the workflow is not already running."""
if self._is_running:
raise RuntimeError("Workflow is already running. Concurrent executions are not allowed.")
self._is_running = True
def _reset_running_flag(self) -> None:
"""Reset the running flag."""
self._is_running = False
def to_dict(self) -> dict[str, Any]:
"""Serialize the workflow definition into a JSON-ready dictionary."""
data: dict[str, Any] = {
@@ -478,6 +478,44 @@ class Workflow(DictConvertible):
"""Get the list of executors in the workflow."""
return list(self.executors.values())
async def _ensure_initial_checkpoint(self) -> None:
"""Capture the in-memory initial checkpoint once for this workflow instance."""
if self._initial_checkpoint is not None:
return
self._initial_checkpoint = await self._runner.capture_checkpoint_object(
metadata={"kind": "initial_in_memory"},
)
async def reset(self) -> None:
"""Reset the workflow instance to its captured initial checkpoint state.
The initial checkpoint is captured in memory once per workflow instance and
is not persisted to external checkpoint storage.
Raises:
WorkflowException: If called while a workflow run is active.
"""
if self._is_run_active():
raise WorkflowException(
"Cannot reset workflow while a run is active. "
"Reset is only allowed between runs when the workflow is idle."
)
# Capture the baseline if it doesn't exist yet. This is idempotent: on a
# normal reset after one or more runs it's a no-op (the snapshot was taken
# before the first run); when reset is the first operation it captures the
# pristine just-built state so the workflow stays runnable.
await self._ensure_initial_checkpoint()
if self._initial_checkpoint is None:
raise WorkflowException("Workflow initial checkpoint is unavailable.")
# Restore runner state, executor snapshots, and runtime bookkeeping from the
# in-memory initial checkpoint.
await self._runner.restore_from_checkpoint_object(self._initial_checkpoint)
self._status = WorkflowRunState.IDLE
async def _run_workflow_with_tracing(
self,
initial_executor_fn: Callable[[], Awaitable[None]] | None = None,
@@ -535,13 +573,12 @@ class Workflow(DictConvertible):
yield in_progress # noqa: RUF070
# Per-run reset for fresh-message runs only. We deliberately
# do NOT clear shared workflow state (`_state.clear()`) or the
# runner context's in-flight messages (`reset_for_new_run()`)
# here - state and pending work persist across `run()` calls
# so that a `WorkflowAgent` can deliver multi-turn input on
# the same instance and have prior turns' context survive.
# Iteration counting and per-run kwargs ARE per-run though,
# so they're reset here.
# do NOT clear shared workflow state or the runner context's
# in-flight messages here - state and pending work persist
# across `run()` calls so that a `WorkflowAgent` can deliver
# multi-turn input on the same instance and have prior turns'
# context survive. Iteration counting and per-run kwargs ARE
# per-run though, so they're reset here.
if not is_continuation:
self._runner.reset_iteration_count()
@@ -564,14 +601,13 @@ class Workflow(DictConvertible):
combined_kwargs["client_kwargs"] = self._resolve_invocation_kwargs(
client_kwargs, "client_kwargs"
)
self._state.set(WORKFLOW_RUN_KWARGS_KEY, combined_kwargs)
self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, combined_kwargs)
elif not is_continuation:
self._state.set(WORKFLOW_RUN_KWARGS_KEY, {})
self._state.commit() # Commit immediately so kwargs are available
self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, {})
self._runner.state.commit() # Commit immediately so kwargs are available
# Set streaming mode (always set explicitly per run since
# reset_for_new_run() no longer runs to clear it).
self._runner_context.set_streaming(streaming)
# Explicitly set streaming mode per run
self._runner.context.set_streaming(streaming)
# Execute initial setup if provided
if initial_executor_fn:
@@ -665,7 +701,7 @@ class Workflow(DictConvertible):
await executor.execute(
message,
[self.__class__.__name__],
self._state,
self._runner.state,
self._runner.context,
trace_contexts=None,
source_span_ids=None,
@@ -745,9 +781,22 @@ class Workflow(DictConvertible):
Raises:
ValueError: If parameter combination is invalid.
"""
# Validate parameters and set running flag eagerly (before any async work)
# Validate parameters first so misuse fails before we touch any run state.
self._validate_run_params(message, responses, checkpoint_id)
self._ensure_not_running()
# Concurrency check: reject a second run synchronously - before constructing
# the ResponseStream or yielding control to the event loop - so a concurrent
# ``run`` call can't slip past the guard while the first call is suspended
# inside its async generator. The ``ResponseStream`` returned below is the
# lock: as long as the caller holds a reference to it, ``self._active_run()``
# resolves to a live object and a new ``run`` is rejected. When the stream is
# fully consumed, ``_run_core``'s finally clears the attribute. When the
# caller drops the stream without iterating, garbage collection invalidates
# the weakref, so a subsequent ``run`` is permitted.
if self._is_run_active():
raise WorkflowException(
"Workflow is already running; concurrent runs are not allowed on the same instance."
)
response_stream = ResponseStream[WorkflowEvent, WorkflowRunResult](
self._run_core(
@@ -760,10 +809,8 @@ class Workflow(DictConvertible):
client_kwargs=client_kwargs,
),
finalizer=functools.partial(self._finalize_events, include_status_events=include_status_events),
cleanup_hooks=[
functools.partial(self._run_cleanup, checkpoint_storage),
],
)
self._active_run = weakref.ref(response_stream)
if stream:
return response_stream
@@ -789,51 +836,69 @@ class Workflow(DictConvertible):
if checkpoint_storage is not None:
self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage)
# Async validation: a fresh-message run is only allowed when the
# runner context has fully drained from any prior run. If it still
# has in-flight executor messages, the prior run didn't complete -
# the caller must either resume from a checkpoint or wait for the
# prior run to drain. (Pending request_info events are intentionally
# NOT blocked here: a follow-up run with message=... is the normal
# way to deliver a response to those pending requests, e.g. via
# WorkflowAgent._process_pending_requests.)
# NOTE: _validate_run_params already enforces that ``message`` is
# mutually exclusive with both ``checkpoint_id`` and ``responses``,
# so we don't need to re-check those here.
if message is not None and await self._runner.context.has_messages():
raise RuntimeError(
"Cannot start a new run with 'message' while in-flight executor "
"messages remain from a prior run. Resume from a checkpoint "
"(checkpoint_id=...) or wait for the prior run to complete. "
"Workflows that need to recover from a mid-run failure must use "
"checkpointing; there is no in-process recovery path."
)
# Capture the weakref instance ``run()`` installed for *this* run. We
# compare by object identity in the finally so a stale finalizer (e.g.
# the caller dropped this stream after partial iteration, then started
# a new run before async-gen finalization throws ``GeneratorExit`` into
# us) does not clobber a successor run's freshly installed weakref.
# ``run()`` runs synchronously and assigns ``self._active_run`` before
# this generator's body is first iterated, so by the time we read it
# here it already points at our own ``ResponseStream``.
my_active_run = self._active_run
initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage)
try:
# Async validation: a fresh-message run is only allowed when the
# runner context has fully drained from any prior run. If it still
# has in-flight executor messages, the prior run didn't complete -
# the caller must either resume from a checkpoint or wait for the
# prior run to drain. (Pending request_info events are intentionally
# NOT blocked here: a follow-up run with message=... is the normal
# way to deliver a response to those pending requests, e.g. via
# WorkflowAgent._process_pending_requests.)
# NOTE: _validate_run_params already enforces that ``message`` is
# mutually exclusive with both ``checkpoint_id`` and ``responses``,
# so we don't need to re-check those here.
if message is not None and await self._runner.context.has_messages():
raise RuntimeError(
"Cannot start a new run with 'message' while in-flight executor "
"messages remain from a prior run. Resume from a checkpoint "
"(checkpoint_id=...) or wait for the prior run to complete. "
"Workflows that need to recover from a mid-run failure must use "
"checkpointing; there is no in-process recovery path."
)
async for event in self._run_workflow_with_tracing(
initial_executor_fn=initial_executor_fn,
is_continuation=(message is None),
streaming=streaming,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
if event.type == "request_info" and event.request_id in (responses or {}):
# Don't yield request_info events for which we have responses to send -
# these are considered "handled". This prevents the caller from seeing
# events for requests they are already responding to.
# This usually happens when responses are provided with a checkpoint
# (restore then send), because the request_info events are stored in the
# checkpoint and would be emitted on restoration by the runner regardless
# of if a response is provided or not.
continue
yield event
await self._ensure_initial_checkpoint()
async def _run_cleanup(self, checkpoint_storage: CheckpointStorage | None) -> None:
"""Cleanup hook called after stream consumption."""
if checkpoint_storage is not None:
self._runner.context.clear_runtime_checkpoint_storage()
self._reset_running_flag()
initial_executor_fn = self._resolve_execution_mode(message, responses, checkpoint_id, checkpoint_storage)
async for event in self._run_workflow_with_tracing(
initial_executor_fn=initial_executor_fn,
is_continuation=(message is None),
streaming=streaming,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
if event.type == "request_info" and event.request_id in (responses or {}):
# Don't yield request_info events for which we have responses to send -
# these are considered "handled". This prevents the caller from seeing
# events for requests they are already responding to.
# This usually happens when responses are provided with a checkpoint
# (restore then send), because the request_info events are stored in the
# checkpoint and would be emitted on restoration by the runner regardless
# of if a response is provided or not.
continue
yield event
finally:
# Clear the active-run weakref so a subsequent ``run()`` is allowed,
# but only if the slot still holds *our* weakref. If the caller
# dropped this stream after partial iteration and a new ``run()``
# already installed its own weakref before our async-gen finalizer
# ran, ``self._active_run`` now points at the successor; clearing
# it would silently break the successor's concurrency guard.
if self._active_run is my_active_run:
self._active_run = None
if checkpoint_storage is not None:
self._runner.context.clear_runtime_checkpoint_storage()
@staticmethod
def _finalize_events(
@@ -935,7 +1000,7 @@ class Workflow(DictConvertible):
async def _send_responses_internal(self, responses: Mapping[str, Any]) -> None:
"""Internal method to validate and send responses to the executors."""
pending_requests = await self._runner_context.get_pending_request_info_events()
pending_requests = await self._runner.context.get_pending_request_info_events()
if not pending_requests:
raise RuntimeError("No pending requests found in workflow context.")
@@ -955,7 +1020,7 @@ class Workflow(DictConvertible):
coerced_responses[request_id] = response
await asyncio.gather(*[
self._runner_context.send_request_info_response(request_id, response)
self._runner.context.send_request_info_response(request_id, response)
for request_id, response in coerced_responses.items()
])
@@ -1151,3 +1216,12 @@ class Workflow(DictConvertible):
context_providers=context_providers,
**kwargs,
)
def _is_run_active(self) -> bool:
"""Check if a workflow run is currently active.
Returns:
True if a run is active, False otherwise.
"""
existing_stream = self._active_run() if self._active_run is not None else None
return existing_stream is not None
@@ -517,6 +517,10 @@ class WorkflowExecutor(Executor):
self._execution_contexts = execution_contexts
self._request_to_execution = request_to_execution
# Reset the sub workflow to its initial state. This must be done before pumping
# the request info events back into the sub workflow.
await self.workflow.reset()
# Add the `request_info_event`s back to the sub workflow.
# This is only a temporary solution to rehydrate the sub workflow with the requests.
# The proper way would be to rehydrate the workflow from a checkpoint on a Workflow
@@ -336,6 +336,97 @@ async def test_workflow_checkpoint_chaining_via_previous_checkpoint_id():
)
async def test_workflow_checkpoint_ancestry_preserved_after_resume():
"""Resuming from a checkpoint must preserve ancestry: future checkpoints chain back to the resumed one."""
from typing_extensions import Never
from agent_framework import WorkflowBuilder, WorkflowContext, handler
from agent_framework._workflows._executor import Executor
class StartExecutor(Executor):
@handler
async def run(self, message: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message(message, target_id="middle")
class MiddleExecutor(Executor):
@handler
async def process(self, message: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message(message + "-processed", target_id="finish")
class FinishExecutor(Executor):
@handler
async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None:
await ctx.yield_output(message + "-done")
storage = InMemoryCheckpointStorage()
def _build_workflow() -> Any:
start = StartExecutor(id="start")
middle = MiddleExecutor(id="middle")
finish = FinishExecutor(id="finish")
return (
WorkflowBuilder(
name="resume-ancestry-test",
max_iterations=10,
start_executor=start,
checkpoint_storage=storage,
)
.add_edge(start, middle)
.add_edge(middle, finish)
.build()
)
# First run: produce an initial chain of checkpoints
workflow = _build_workflow()
workflow_name = workflow.name
_ = [event async for event in workflow.run("hello", stream=True)]
initial_checkpoints = sorted(await storage.list_checkpoints(workflow_name=workflow_name), key=lambda c: c.timestamp)
assert len(initial_checkpoints) >= 3, (
f"Need at least 3 initial checkpoints to pick a middle one, got {len(initial_checkpoints)}"
)
initial_ids = {cp.checkpoint_id for cp in initial_checkpoints}
# Pick an intermediate checkpoint to resume from (not the first, not the last)
resume_from = initial_checkpoints[len(initial_checkpoints) // 2]
# Resume on a fresh workflow instance (same graph signature) and run to completion
resumed_workflow = _build_workflow()
assert resumed_workflow.name == workflow_name
_ = [event async for event in resumed_workflow.run(checkpoint_id=resume_from.checkpoint_id, stream=True)]
# Inspect new checkpoints created after resuming
all_checkpoints = sorted(await storage.list_checkpoints(workflow_name=workflow_name), key=lambda c: c.timestamp)
new_checkpoints = [cp for cp in all_checkpoints if cp.checkpoint_id not in initial_ids]
assert new_checkpoints, "Resuming from an intermediate checkpoint should produce new checkpoints"
# The very first checkpoint created after resuming must chain back to the resumed checkpoint
assert new_checkpoints[0].previous_checkpoint_id == resume_from.checkpoint_id, (
"First post-resume checkpoint must chain to the checkpoint that was resumed from; "
f"got previous_checkpoint_id={new_checkpoints[0].previous_checkpoint_id!r}, "
f"expected {resume_from.checkpoint_id!r}"
)
# Subsequent post-resume checkpoints must continue chaining
for i in range(1, len(new_checkpoints)):
assert new_checkpoints[i].previous_checkpoint_id == new_checkpoints[i - 1].checkpoint_id, (
f"Post-resume checkpoint {i} should chain to checkpoint {i - 1}"
)
# Walking the chain backwards from the most recent checkpoint must reach the original root
# without breaks (i.e. the full ancestry across the resume boundary is intact).
checkpoints_by_id = {cp.checkpoint_id: cp for cp in all_checkpoints}
chain: list[str] = []
cursor: str | None = new_checkpoints[-1].checkpoint_id
while cursor is not None:
chain.append(cursor)
cursor = checkpoints_by_id[cursor].previous_checkpoint_id
# Chain must include the resumed-from checkpoint and terminate at the original root
assert resume_from.checkpoint_id in chain
assert chain[-1] == initial_checkpoints[0].checkpoint_id
assert checkpoints_by_id[chain[-1]].previous_checkpoint_id is None
async def test_memory_checkpoint_storage_roundtrip_json_native_types():
"""Test that JSON-native types (str, int, float, bool, None) roundtrip correctly."""
storage = InMemoryCheckpointStorage()
@@ -17,7 +17,6 @@ from agent_framework import (
WorkflowContext,
WorkflowConvergenceException,
WorkflowEvent,
WorkflowRunnerException,
WorkflowRunState,
handler,
)
@@ -305,40 +304,62 @@ async def test_fanout_edge_runner_delivers_to_multiple_targets_concurrently() ->
assert probe_target.call_count == 1
async def test_runner_already_running():
"""Test that running the runner while it is already running raises an error."""
async def test_runner_run_until_convergence_runs_sequentially():
"""run_until_convergence can be invoked back-to-back on the same Runner.
The Runner itself does not enforce concurrency; that responsibility lives on
:class:`Workflow`. This test simply confirms the Runner is reusable across
sequential runs.
"""
runner = _make_runner()
async for _ in runner.run_until_convergence():
pass
async for _ in runner.run_until_convergence():
pass
def _make_runner() -> Runner:
"""Build a minimal runner for runner-level tests."""
return Runner(
[],
{},
State(),
InProcRunnerContext(),
"test_name",
graph_signature_hash="test_hash",
)
async def test_runner_accepts_new_run_after_previous_failure():
"""A failed run must not leave the Runner unable to start a new run.
After the first run raises, ``run_until_convergence()`` must be callable
again and not surface any lifecycle-related rejection.
"""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
# Create a loop
edges = [
SingleEdgeGroup(executor_a.id, executor_b.id),
SingleEdgeGroup(executor_b.id, executor_a.id),
]
executors: dict[str, Executor] = {
executor_a.id: executor_a,
executor_b.id: executor_b,
}
executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
state = State()
ctx = InProcRunnerContext()
runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash", max_iterations=2)
runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
await executor_a.execute(MockMessage(data=0), ["START"], state, ctx)
await executor_a.execute(
MockMessage(data=0),
["START"], # source_executor_ids
state, # state
ctx, # runner_context
)
with pytest.raises(WorkflowConvergenceException):
async for _ in runner.run_until_convergence():
pass
with pytest.raises(WorkflowRunnerException, match="Runner is already running."):
async def _run():
async for _ in runner.run_until_convergence():
pass
await asyncio.gather(_run(), _run())
# A second run on the same Runner must not be blocked by stale lifecycle
# state from the failed run.
try:
async for _ in runner.run_until_convergence():
pass
except Exception as exc:
assert "Runner is already running" not in str(exc), "Runner stayed locked after a failed run"
async def test_runner_emits_runner_completion_for_agent_response_without_targets():
@@ -862,7 +883,13 @@ async def test_runner_checkpoint_with_resumed_flag():
state = State()
runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
runner._mark_resumed(5) # pyright: ignore[reportPrivateUsage]
resumed_checkpoint = WorkflowCheckpoint(
checkpoint_id="resumed-cp",
workflow_name="test_name",
graph_signature_hash="test_hash",
iteration_count=5,
)
runner._mark_resumed(resumed_checkpoint) # pyright: ignore[reportPrivateUsage]
# Add a message to trigger the checkpoint creation path
await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id="START"))
@@ -882,6 +909,86 @@ async def test_runner_checkpoint_with_resumed_flag():
assert runner._resumed_from_checkpoint is False # pyright: ignore[reportPrivateUsage]
async def test_runner_mark_resumed_sets_previous_checkpoint_id():
"""_mark_resumed must populate _previous_checkpoint_id so future checkpoints chain back to the resume point."""
runner = Runner(
[],
{},
State(),
InProcRunnerContext(),
"test_name",
graph_signature_hash="test_hash",
)
# Pre-condition: nothing to chain back to
assert runner._previous_checkpoint_id is None # pyright: ignore[reportPrivateUsage]
resumed_checkpoint = WorkflowCheckpoint(
checkpoint_id="resumed-cp-id",
workflow_name="test_name",
graph_signature_hash="test_hash",
iteration_count=3,
)
runner._mark_resumed(resumed_checkpoint) # pyright: ignore[reportPrivateUsage]
assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage]
assert runner._iteration == 3 # pyright: ignore[reportPrivateUsage]
assert runner._previous_checkpoint_id == "resumed-cp-id" # pyright: ignore[reportPrivateUsage]
async def test_runner_post_resume_checkpoint_chains_to_resumed_checkpoint():
"""After resuming, the next checkpoint created must reference the resumed checkpoint as its parent."""
storage = InMemoryCheckpointStorage()
ctx = CheckpointingContext(storage)
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
edges = [
SingleEdgeGroup(executor_a.id, executor_b.id),
SingleEdgeGroup(executor_b.id, executor_a.id),
]
executors: dict[str, Executor] = {
executor_a.id: executor_a,
executor_b.id: executor_b,
}
state = State()
runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
# Simulate having resumed from a prior checkpoint
resumed_checkpoint = WorkflowCheckpoint(
checkpoint_id="parent-checkpoint-id",
workflow_name="test_name",
graph_signature_hash="test_hash",
iteration_count=1,
)
runner._mark_resumed(resumed_checkpoint) # pyright: ignore[reportPrivateUsage]
# Seed a message so the runner has work to do (and creates checkpoints at superstep boundaries)
await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id=executor_a.id))
async for _ in runner.run_until_convergence():
pass
# Find the first checkpoint created after the resume point (across all workflows tracked by storage)
new_checkpoints = sorted(
await storage.list_checkpoints(workflow_name="test_name"),
key=lambda c: c.timestamp,
)
assert new_checkpoints, "Resuming and running should produce at least one new checkpoint"
# The first new checkpoint must chain to the resumed-from checkpoint, not to None
assert new_checkpoints[0].previous_checkpoint_id == "parent-checkpoint-id", (
"First post-resume checkpoint must chain to the resumed checkpoint id; "
f"got {new_checkpoints[0].previous_checkpoint_id!r}"
)
# Subsequent post-resume checkpoints continue the chain
for i in range(1, len(new_checkpoints)):
assert new_checkpoints[i].previous_checkpoint_id == new_checkpoints[i - 1].checkpoint_id
class ExecutorThatFailsWithEvents(Executor):
"""An executor that emits events and then raises an exception after receiving messages."""
@@ -0,0 +1,72 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for `InProcRunnerContext`."""
import pytest
from agent_framework import (
InProcRunnerContext,
WorkflowEvent,
WorkflowMessage,
)
def _make_request_info_event(request_id: str, source_executor_id: str = "executor") -> WorkflowEvent[str]:
return WorkflowEvent.request_info(
request_id=request_id,
source_executor_id=source_executor_id,
request_data="please respond",
response_type=str,
)
class TestInProcRunnerContextResetForNewRun:
"""Verify `reset_for_new_run` clears per-run state, including pending request_info events."""
async def test_reset_clears_pending_request_info_events(self) -> None:
ctx = InProcRunnerContext()
await ctx.add_request_info_event(_make_request_info_event("req-1"))
await ctx.add_request_info_event(_make_request_info_event("req-2"))
assert set((await ctx.get_pending_request_info_events()).keys()) == {"req-1", "req-2"}
ctx.reset_for_new_run()
assert await ctx.get_pending_request_info_events() == {}
async def test_reset_clears_pending_request_info_events_when_already_empty(self) -> None:
ctx = InProcRunnerContext()
assert await ctx.get_pending_request_info_events() == {}
ctx.reset_for_new_run()
assert await ctx.get_pending_request_info_events() == {}
async def test_reset_after_pending_event_blocks_response_correlation(self) -> None:
"""After `reset_for_new_run`, prior request ids must no longer correlate to a response."""
ctx = InProcRunnerContext()
await ctx.add_request_info_event(_make_request_info_event("req-1"))
ctx.reset_for_new_run()
with pytest.raises(ValueError, match="No pending request found for request_id: req-1"):
await ctx.send_request_info_response("req-1", "answer")
async def test_reset_clears_messages_events_and_streaming_flag(self) -> None:
"""Sanity-check the other state `reset_for_new_run` is documented to clear."""
ctx = InProcRunnerContext()
await ctx.send_message(WorkflowMessage(data="hello", source_id="executor"))
await ctx.add_event(WorkflowEvent("status", data="running"))
ctx.set_streaming(True)
assert await ctx.has_messages() is True
assert await ctx.has_events() is True
assert ctx.is_streaming() is True
ctx.reset_for_new_run()
assert await ctx.has_messages() is False
assert await ctx.has_events() is False
assert ctx.is_streaming() is False
@@ -1,6 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import gc
import tempfile
from collections.abc import AsyncIterable, Awaitable, Sequence
from dataclasses import dataclass, field
@@ -26,6 +27,7 @@ from agent_framework import (
WorkflowContext,
WorkflowConvergenceException,
WorkflowEvent,
WorkflowException,
WorkflowMessage,
WorkflowRunState,
handler,
@@ -759,8 +761,7 @@ async def test_workflow_concurrent_execution_prevention():
# Try to start a second concurrent execution - this should fail
with pytest.raises(
RuntimeError,
match="Workflow is already running. Concurrent executions are not allowed.",
WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
):
await workflow.run(NumberMessage(data=0))
@@ -795,8 +796,7 @@ async def test_workflow_concurrent_execution_prevention_streaming():
# Try to start a second concurrent execution - this should fail
with pytest.raises(
RuntimeError,
match="Workflow is already running. Concurrent executions are not allowed.",
WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
):
await workflow.run(NumberMessage(data=0))
@@ -828,14 +828,12 @@ async def test_workflow_concurrent_execution_prevention_mixed_methods():
# Try different execution methods - all should fail
with pytest.raises(
RuntimeError,
match="Workflow is already running. Concurrent executions are not allowed.",
WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
):
await workflow.run(NumberMessage(data=0))
with pytest.raises(
RuntimeError,
match="Workflow is already running. Concurrent executions are not allowed.",
WorkflowException, match="Workflow is already running; concurrent runs are not allowed on the same instance."
):
async for _ in workflow.run(NumberMessage(data=0), stream=True):
break
@@ -848,6 +846,154 @@ async def test_workflow_concurrent_execution_prevention_mixed_methods():
assert result.get_final_state() == WorkflowRunState.IDLE
async def test_workflow_sequential_runs_after_completion() -> None:
"""A completed run must release the runner so the next ``run`` succeeds.
This is the happy-path counterpart to the concurrent-run guard tests:
those tests verify that a *concurrent* run is rejected, but they do not
verify that the lock is actually released afterwards. This test
exercises that release path explicitly across the three call shapes
(non-streaming, streaming-iterated, streaming-via-get_final_response)
and across multiple consecutive turns to catch lock leaks.
"""
executor = IncrementExecutor(id="seq_executor", limit=3, increment=1)
workflow = WorkflowBuilder(start_executor=executor).build()
# Non-streaming -> non-streaming
r1 = await workflow.run(NumberMessage(data=0))
assert r1.get_final_state() == WorkflowRunState.IDLE
r2 = await workflow.run(NumberMessage(data=0))
assert r2.get_final_state() == WorkflowRunState.IDLE
# Non-streaming -> streaming-iterated
stream_events: list[WorkflowEvent] = []
async for event in workflow.run(NumberMessage(data=0), stream=True):
stream_events.append(event)
assert any(e.type == "status" and e.state == WorkflowRunState.IDLE for e in stream_events)
# Streaming -> streaming via get_final_response (no manual iteration)
r3 = await workflow.run(NumberMessage(data=0), stream=True).get_final_response()
assert r3.get_final_state() == WorkflowRunState.IDLE
# Streaming -> non-streaming (back to the start)
r4 = await workflow.run(NumberMessage(data=0))
assert r4.get_final_state() == WorkflowRunState.IDLE
async def test_workflow_unconsumed_stream_releases_run_lock() -> None:
"""An unconsumed stream must not leak the run lock.
``Workflow.run`` reserves the runner *synchronously* so that concurrent
callers are rejected immediately. The reservation is normally released
by ``_run_core``'s ``finally`` once the stream is iterated. If the
caller never iterates the stream, a GC-time finalizer must release the
reservation instead - otherwise every subsequent ``Workflow.run`` call
on this instance would fail with the concurrent-run error.
"""
executor = IncrementExecutor(id="unconsumed_stream_exec", limit=3, increment=1)
workflow = WorkflowBuilder(start_executor=executor).build()
# Build a stream and immediately drop it without iterating.
stream = workflow.run(NumberMessage(data=0), stream=True)
assert stream is not None # silence unused-variable warnings; stream is GC'd below
del stream
gc.collect()
# Yield to the event loop so any scheduled finalizer work can run.
await asyncio.sleep(0)
# The runner should be back to IDLE; a fresh run must succeed.
result = await workflow.run(NumberMessage(data=0))
assert result.get_final_state() == WorkflowRunState.IDLE
async def test_workflow_unawaited_run_coroutine_releases_run_lock() -> None:
"""An un-awaited non-streaming ``run()`` coroutine must also not leak the lock.
``Workflow.run`` (non-streaming) returns a coroutine produced by
``ResponseStream.get_final_response``. The underlying ResponseStream is
held alive by that coroutine, so dropping the coroutine without
awaiting it must still release the reservation via the same GC-time
fallback used for unconsumed streams.
"""
executor = IncrementExecutor(id="unawaited_run_exec", limit=3, increment=1)
workflow = WorkflowBuilder(start_executor=executor).build()
coro = workflow.run(NumberMessage(data=0))
# Closing suppresses the "coroutine was never awaited" warning. We cast to
# ``Any`` because the typed return is ``Awaitable[...]``; in practice it is
# a coroutine that exposes ``close``.
cast(Any, coro).close()
del coro
gc.collect()
await asyncio.sleep(0)
result = await workflow.run(NumberMessage(data=0))
assert result.get_final_state() == WorkflowRunState.IDLE
async def test_workflow_partial_stream_does_not_clobber_successor_active_run() -> None:
"""A stale ``_run_core`` finalizer must not clear a successor's run lock.
Repro for the GC-finalizer race the user reported:
1. Start stream A and consume one event so its body is suspended at a
``yield``. Its ``finally`` is now armed and will run when the
generator is closed.
2. Drop stream A and ``gc.collect``. The ``_active_run`` weakref's
referent is gone, so a subsequent ``run()`` will pass the
concurrency guard - but stream A's async-gen finalizer hasn't
actually executed yet (``aclose`` is scheduled on the loop).
3. Synchronously start stream B; ``run()`` installs a fresh weakref
in ``_active_run``.
4. Yield to the loop so stream A's stale ``finally`` runs. Without
the identity check it unconditionally writes
``self._active_run = None``, silently disabling the concurrency
guard for stream B.
"""
executor = IncrementExecutor(id="stale_finalizer_exec", limit=100, increment=1)
workflow = WorkflowBuilder(start_executor=executor).build()
# Step 1: drive stream A's body until it's suspended at its first yield.
stream_a = workflow.run(NumberMessage(data=0), stream=True)
aiter_a = stream_a.__aiter__()
await aiter_a.__anext__()
# Step 2: drop stream A; GC invalidates the weakref and schedules
# async-gen close, but does not run the close inline.
del stream_a
del aiter_a
gc.collect()
# Step 3: synchronously start stream B *before* yielding to the loop,
# so the stale ``aclose`` for stream A hasn't fired yet.
stream_b = workflow.run(NumberMessage(data=0), stream=True)
ref_b = workflow._active_run # type: ignore[attr-defined]
assert ref_b is not None and ref_b() is stream_b
# Step 4: yield enough times for stream A's scheduled aclose to drive
# its body through ``GeneratorExit`` and into its ``finally``.
for _ in range(5):
await asyncio.sleep(0)
# With the fix, stream B's reservation is still in place. Without it,
# ``_active_run`` was clobbered to ``None`` and a concurrent run would
# be (incorrectly) accepted.
assert workflow._active_run is ref_b # type: ignore[attr-defined]
with pytest.raises(
WorkflowException,
match="Workflow is already running; concurrent runs are not allowed on the same instance.",
):
await workflow.run(NumberMessage(data=0))
# Tear down stream B without iterating it (its body never started, so
# closing it is a no-op for workflow state).
del stream_b
del ref_b
gc.collect()
await asyncio.sleep(0)
class _StreamingTestAgent(BaseAgent):
"""Test agent that supports both streaming and non-streaming modes."""
@@ -1269,3 +1415,85 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None
# endregion
# region Workflow.reset
class CounterStateExecutor(Executor):
"""Executor with local mutable state used to verify checkpoint-based reset."""
def __init__(self, id: str) -> None:
super().__init__(id=id)
self.counter = 0
@handler
async def handle(self, message: str, ctx: WorkflowContext[str, int]) -> None:
self.counter += 1
await ctx.yield_output(self.counter)
async def on_checkpoint_save(self) -> dict[str, Any]:
return {"counter": self.counter}
async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
self.counter = int(state.get("counter", 0))
class TestWorkflowReset:
"""Tests for :meth:`Workflow.reset`."""
async def test_reset_restores_initial_shared_state(self) -> None:
"""Reset clears accumulated workflow state back to the initial baseline."""
executor = StateTrackingExecutor(id="state_executor")
workflow = WorkflowBuilder(start_executor=executor).add_edge(executor, executor).build()
result1 = await workflow.run(StateTrackingMessage(data="message1", run_id="run1"))
assert result1.get_outputs()[0] == ["run1:message1"]
result2 = await workflow.run(StateTrackingMessage(data="message2", run_id="run2"))
assert result2.get_outputs()[0] == ["run1:message1", "run2:message2"]
await workflow.reset()
result3 = await workflow.run(StateTrackingMessage(data="message3", run_id="run3"))
assert result3.get_outputs()[0] == ["run3:message3"]
async def test_reset_restores_executor_checkpoint_state(self) -> None:
"""Reset restores per-executor local state captured in the initial checkpoint."""
executor = CounterStateExecutor(id="counter_executor")
workflow = WorkflowBuilder(start_executor=executor).add_edge(executor, executor).build()
result1 = await workflow.run("one")
assert result1.get_outputs() == [1]
result2 = await workflow.run("two")
assert result2.get_outputs() == [2]
await workflow.reset()
result3 = await workflow.run("three")
assert result3.get_outputs() == [1]
async def test_reset_before_first_run_is_allowed(self, simple_executor: Executor) -> None:
"""Reset can be called before the first run and leaves workflow runnable."""
workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
await workflow.reset()
result = await workflow.run("hello")
assert result.get_final_state() == WorkflowRunState.IDLE
async def test_reset_raises_while_run_active(self, simple_executor: Executor) -> None:
"""Reset must reject while a workflow run is active."""
workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
active_stream = workflow.run(WorkflowMessage(data="hi", source_id="test"), stream=True)
try:
with pytest.raises(WorkflowException, match="Cannot reset workflow while a run is active"):
await workflow.reset()
finally:
async for _ in active_stream:
pass
# endregion
@@ -90,7 +90,7 @@ async def _run(yaml_def: dict[str, Any], handler: HttpRequestHandler) -> Any:
def _state(workflow: Any, events: Any) -> dict[str, Any]:
"""Read declarative state out of the workflow after run completes."""
return workflow._state.get(DECLARATIVE_STATE_KEY) or {}
return workflow._runner.state.get(DECLARATIVE_STATE_KEY) or {}
# Helper used by parametrised path tests
@@ -151,7 +151,7 @@ class TestSuccessPath:
workflow = factory.create_workflow_from_definition(_yaml(_action(method="GET", response="Local.Result")))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == {"key": "value", "number": 42}
assert handler.last_info is not None
assert handler.last_info.method == "GET"
@@ -164,7 +164,7 @@ class TestSuccessPath:
workflow = factory.create_workflow_from_definition(_yaml(_action(response="Local.Result")))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == "not-json content"
@pytest.mark.asyncio
@@ -174,7 +174,7 @@ class TestSuccessPath:
workflow = factory.create_workflow_from_definition(_yaml(_action(response="Local.Result")))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] is None
@pytest.mark.asyncio
@@ -184,7 +184,7 @@ class TestSuccessPath:
workflow = factory.create_workflow_from_definition(_yaml(_action(response={"path": "Local.Result"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == {"x": 1}
@pytest.mark.asyncio
@@ -517,7 +517,7 @@ class TestResponseHeaders:
factory = WorkflowFactory(http_request_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(response_headers="Local.H")))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
h = decl["Local"]["H"]
assert h["Content-Type"] == "application/json"
assert h["Set-Cookie"] == "a=1,b=2"
@@ -528,7 +528,7 @@ class TestResponseHeaders:
factory = WorkflowFactory(http_request_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(response_headers="Local.H")))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["H"] is None
@pytest.mark.asyncio
@@ -538,7 +538,7 @@ class TestResponseHeaders:
workflow = factory.create_workflow_from_definition(_yaml(_action(response_headers="Local.H")))
with pytest.raises(DeclarativeActionError):
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["H"] == {"X-Trace": "abc"}
@@ -559,7 +559,7 @@ class TestConversationAppend:
)
)
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
conv = decl["System"]["conversations"].get("conv-test-1")
assert conv is not None
assert len(conv["messages"]) == 1
@@ -570,7 +570,7 @@ class TestConversationAppend:
factory = WorkflowFactory(http_request_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(response="Local.Result", conversation_id="")))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
# Auto-init creates an entry for the System.ConversationId conversation,
# but it should NOT have HTTP-appended messages from us.
for _cid, conv in decl["System"]["conversations"].items():
@@ -582,7 +582,7 @@ class TestConversationAppend:
factory = WorkflowFactory(http_request_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(conversation_id="conv-test-1")))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
# No conversation entry should have been created either.
assert "conv-test-1" not in decl["System"]["conversations"]
@@ -73,8 +73,8 @@ async def test_http_request_yaml_roundtrip() -> None:
workflow = factory.create_workflow_from_yaml_path(FIXTURE_PATH)
await workflow.run({})
decl: dict[str, Any] = workflow._state.get(DECLARATIVE_STATE_KEY) or {}
local: dict[str, Any] = decl.get("Local") or {}
decl: dict[str, Any] = workflow._runner.state.get(DECLARATIVE_STATE_KEY) or {}
local = decl.get("Local") or {}
assert local.get("RepoOwner") == "dotnet"
repo_info = local.get("RepoInfo")
@@ -244,7 +244,7 @@ class TestOutput:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == [{"k": "v", "n": 1}]
@pytest.mark.asyncio
@@ -253,7 +253,7 @@ class TestOutput:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == ["plain text not json"]
@pytest.mark.asyncio
@@ -262,7 +262,7 @@ class TestOutput:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"messages": "Local.Messages"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
msg = decl["Local"]["Messages"]
# Single Tool-role message containing both contents (parity with .NET).
assert isinstance(msg, Message)
@@ -276,7 +276,7 @@ class TestOutput:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == ["https://example.com/file.txt"]
@pytest.mark.asyncio
@@ -285,7 +285,7 @@ class TestOutput:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": {"path": "Local.Result"}})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == ["ok"]
@@ -306,7 +306,7 @@ class TestConversation:
)
)
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
conv = decl["System"]["conversations"]["conv-42"]
msgs = conv["messages"] if isinstance(conv, dict) else conv.messages
assert len(msgs) == 1
@@ -328,7 +328,7 @@ class TestConversation:
)
)
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
# Empty conversation id must not produce a `""` entry under System.conversations.
conversations = decl.get("System", {}).get("conversations", {})
assert "" not in conversations
@@ -529,7 +529,7 @@ class TestErrorHandling:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == "Error: server down"
@pytest.mark.asyncio
@@ -538,7 +538,7 @@ class TestErrorHandling:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert decl["Local"]["Result"] == "Error: invalid arguments"
@pytest.mark.asyncio
@@ -547,7 +547,7 @@ class TestErrorHandling:
factory = WorkflowFactory(mcp_tool_handler=handler)
workflow = factory.create_workflow_from_definition(_yaml(_action(output={"result": "Local.Result"})))
await workflow.run({})
decl = workflow._state.get(DECLARATIVE_STATE_KEY)
decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
result = decl["Local"]["Result"]
assert isinstance(result, str)
assert result.startswith("Error:")
@@ -291,11 +291,11 @@ actions:
# Stamp a marker into the declarative state between turns. The
# continuation branch must preserve it; a state-clearing run would
# wipe ``DECLARATIVE_STATE_KEY`` and force re-initialization.
state_data = workflow._state.get(DECLARATIVE_STATE_KEY)
state_data = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert isinstance(state_data, dict), "Expected declarative state to be initialized after turn 1"
state_data["Local"] = {"persisted_marker": "kept-from-turn-1"}
workflow._state.set(DECLARATIVE_STATE_KEY, state_data)
workflow._state.commit()
workflow._runner.state.set(DECLARATIVE_STATE_KEY, state_data)
workflow._runner.state.commit()
second = await agent.run("turn-2-msg")
assert second.text == "turn-2-msg", (
@@ -305,7 +305,7 @@ actions:
# The continuation branch in ``_ensure_state_initialized`` must:
# 1. preserve the cross-turn marker we stamped above
# 2. refresh Inputs.input and System.LastMessage* to the new turn
post_state = workflow._state.get(DECLARATIVE_STATE_KEY)
post_state = workflow._runner.state.get(DECLARATIVE_STATE_KEY)
assert isinstance(post_state, dict), "declarative state vanished between turns"
local = post_state.get("Local", {})
assert local.get("persisted_marker") == "kept-from-turn-1", (
@@ -386,7 +386,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
)
self._is_workflow_agent = False
self._checkpoint_storage_path = None
if isinstance(agent, WorkflowAgent):
if agent.workflow._runner_context.has_checkpointing(): # pyright: ignore[reportPrivateUsage]
raise RuntimeError(
@@ -580,8 +579,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
# The following should never happen due to the checks above.
# This is for type safety and defensive programming.
if self._checkpoint_storage_path is None:
raise RuntimeError("Checkpoint storage path is not configured for workflow agent.")
if not isinstance(self._agent, WorkflowAgent):
raise RuntimeError("Agent is not a workflow agent.")
@@ -599,43 +596,27 @@ class ResponsesHostServer(ResponsesAgentServerHost):
# the only place that state lives is the workflow checkpoint, so
# on every turn we restore the latest checkpoint and feed the new
# input back into the start executor as a continuation rather than
# a fresh run.
# a fresh run. If no conversation_id or previous_response_id is
# supplied (or no checkpoint exists for that context), reset the
# workflow to its in-memory initial baseline to avoid context bleed
# between requests.
latest_checkpoint_id: str | None = None
restore_storage: FileCheckpointStorage | None = None
if context_id is not None:
restore_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, context_id)
latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name)
context_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, context_id)
latest_checkpoint = await context_storage.get_latest(workflow_name=self._agent.workflow.name)
if latest_checkpoint is not None:
latest_checkpoint_id = latest_checkpoint.checkpoint_id
restore_storage = context_storage
# Storage that will receive checkpoints written during this turn.
# When the caller chains with previous_response_id, the next turn
# will reference the current response_id as its previous_response_id,
# so new checkpoints must land under the current response_id (or the
# conversation_id when set). When conversation_id is set, this
# matches restore_storage; when only previous_response_id was
# supplied, restore_storage points at the *prior* response's
# directory and write_storage points at the *current* response's.
write_context_id = context.conversation_id or context.response_id
write_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, write_context_id)
# Multi-turn pattern: when we have a prior checkpoint, restore it
# first (drive the workflow back to idle with prior state intact),
# then make a separate call that delivers the new user input. This
# depends on Workflow.run preserving shared state across calls. The
# restore-only call may yield events from any pending in-flight
# work in the checkpoint; we consume those internally here so they
# don't surface to the response stream as duplicates.
#
# If the restored checkpoint had pending request_info events, the
# restore-only call replays them through
# ``WorkflowAgent._convert_workflow_event_to_agent_response_updates``
# and populates ``self._agent.pending_requests``. That is the correct
# state: those requests are genuinely outstanding, and the next
# ``run(input_messages, ...)`` call may contain ``function_call_output``
# items (carried as FunctionResult/FunctionApprovalResponse content)
# that fulfill them via :meth:`WorkflowAgent._process_pending_requests`.
if latest_checkpoint_id is not None:
# Restore the workflow to the latest checkpoint and run it with the
# new input. Events (including request info events) will not be emitted
# during restoration (in streaming) or after restoration (in non-streaming)
# since we assume the client had already seen those events and we don't want
# to emit duplicates.
if latest_checkpoint_id is None or restore_storage is None:
await self._agent.workflow.reset()
else:
if is_streaming_request:
async for _ in self._agent.run(
stream=True,
@@ -650,6 +631,17 @@ class ResponsesHostServer(ResponsesAgentServerHost):
checkpoint_storage=restore_storage,
)
# Storage that will receive checkpoints written during this turn.
# When the caller chains with previous_response_id, the next turn
# will reference the current response_id as its previous_response_id,
# so new checkpoints must land under the current response_id (or the
# conversation_id when set). When conversation_id is set, this
# matches restore_storage; when only previous_response_id was
# supplied, restore_storage points at the *prior* response's
# directory and write_storage points at the *current* response's.
write_context_id = context.conversation_id or context.response_id
write_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, write_context_id)
if not is_streaming_request:
# Run the agent in non-streaming mode with the new user input.
response = await self._agent.run(
@@ -3062,6 +3062,7 @@ class TestCheckpointContextPathValidation:
agent.workflow = MagicMock()
agent.workflow.name = "wf"
agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
agent.workflow.reset = AsyncMock()
agent.run = AsyncMock(
side_effect=[
AgentResponse(messages=[]),
@@ -3092,6 +3093,136 @@ class TestCheckpointContextPathValidation:
assert new_turn_messages[0].text == "next turn"
assert new_turn_call.kwargs["checkpoint_storage"].storage_path == (root / response_id).resolve()
async def test_handle_inner_workflow_resets_when_no_context_id(self, tmp_path: Any) -> None:
"""When no context id is supplied, the workflow resets to its initial in-memory state."""
from agent_framework import WorkflowAgent
from azure.ai.agentserver.responses import ResponseContext
from azure.ai.agentserver.responses.models import CreateResponse, ItemMessage
response_id = "resp_current"
root = tmp_path / "root"
root.mkdir()
agent = MagicMock(spec=WorkflowAgent)
agent.id = "wf-agent"
agent.name = "wf"
agent.description = ""
agent.context_providers = []
agent.workflow = MagicMock()
agent.workflow.name = "wf"
agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
agent.workflow.reset = AsyncMock()
agent.run = AsyncMock(
return_value=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ok")])])
)
server = ResponsesHostServer(agent, store=InMemoryResponseProvider())
server._checkpoint_storage_path = str(root) # pyright: ignore[reportPrivateUsage]
# No previous_response_id and no conversation_id.
request = CreateResponse(model="m", input="hi")
context = ResponseContext(response_id=response_id, mode_flags=MagicMock())
input_item = ItemMessage({"type": "message", "role": "user", "content": "fresh turn"})
with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])):
async for _ in server._handle_inner_workflow(request, context): # pyright: ignore[reportPrivateUsage]
pass
# No checkpoint restore is attempted; workflow resets in memory.
assert agent.workflow.reset.await_count == 1
assert agent.run.call_count == 1
# The single run() call delivers the new input; checkpoints land under response_id
# (the write-sink directory keyed by the current response id).
new_turn_call = agent.run.call_args_list[0]
new_turn_messages = new_turn_call.args[0]
assert len(new_turn_messages) == 1
assert new_turn_messages[0].text == "fresh turn"
assert new_turn_call.kwargs["checkpoint_storage"].storage_path == (root / response_id).resolve()
async def test_handle_inner_workflow_resets_each_request_without_context_id(self, tmp_path: Any) -> None:
"""Requests without context ids reset workflow state per request."""
from agent_framework import WorkflowAgent
from azure.ai.agentserver.responses import ResponseContext
from azure.ai.agentserver.responses.models import CreateResponse, ItemMessage
root = tmp_path / "root"
root.mkdir()
agent = MagicMock(spec=WorkflowAgent)
agent.id = "wf-agent"
agent.name = "wf"
agent.description = ""
agent.context_providers = []
agent.workflow = MagicMock()
agent.workflow.name = "wf"
agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
agent.workflow.reset = AsyncMock()
# Two run() calls total: one new turn per request.
agent.run = AsyncMock(return_value=AgentResponse(messages=[]))
server = ResponsesHostServer(agent, store=InMemoryResponseProvider())
server._checkpoint_storage_path = str(root) # pyright: ignore[reportPrivateUsage]
request1 = CreateResponse(model="m", input="hi")
context1 = ResponseContext(response_id="resp_first", mode_flags=MagicMock())
request2 = CreateResponse(model="m", input="hi again")
context2 = ResponseContext(response_id="resp_second", mode_flags=MagicMock())
input_item = ItemMessage({"type": "message", "role": "user", "content": "turn"})
with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])):
async for _ in server._handle_inner_workflow(request1, context1): # pyright: ignore[reportPrivateUsage]
pass
async for _ in server._handle_inner_workflow(request2, context2): # pyright: ignore[reportPrivateUsage]
pass
assert agent.workflow.reset.await_count == 2
assert agent.run.call_count == 2
async def test_handle_inner_workflow_resets_when_context_dir_is_empty(self, tmp_path: Any) -> None:
"""When previous_response_id has no checkpoint, workflow resets instead of restoring."""
from agent_framework import WorkflowAgent
from azure.ai.agentserver.responses import ResponseContext
from azure.ai.agentserver.responses.models import CreateResponse, ItemMessage
previous_response_id = "resp_previous"
response_id = "resp_current"
root = tmp_path / "root"
root.mkdir()
# The per-context storage exists but contains no checkpoints.
(root / previous_response_id).mkdir()
agent = MagicMock(spec=WorkflowAgent)
agent.id = "wf-agent"
agent.name = "wf"
agent.description = ""
agent.context_providers = []
agent.workflow = MagicMock()
agent.workflow.name = "wf"
agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
agent.workflow.reset = AsyncMock()
agent.run = AsyncMock(
return_value=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ok")])])
)
server = ResponsesHostServer(agent, store=InMemoryResponseProvider())
server._checkpoint_storage_path = str(root) # pyright: ignore[reportPrivateUsage]
request = CreateResponse(model="m", input="hi", previous_response_id=previous_response_id)
context = ResponseContext(
response_id=response_id, previous_response_id=previous_response_id, mode_flags=MagicMock()
)
input_item = ItemMessage({"type": "message", "role": "user", "content": "next turn"})
with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])):
async for _ in server._handle_inner_workflow(request, context): # pyright: ignore[reportPrivateUsage]
pass
assert agent.workflow.reset.await_count == 1
assert agent.run.call_count == 1
# The new turn writes checkpoints under the current response id.
new_turn_call = agent.run.call_args_list[0]
assert new_turn_call.kwargs["checkpoint_storage"].storage_path == (root / response_id).resolve()
@pytest.mark.parametrize(
"bad_id",
[
@@ -3185,6 +3316,8 @@ class TestCheckpointContextPathValidation:
agent.workflow = MagicMock()
agent.workflow.name = "wf"
agent.workflow._runner_context.has_checkpointing = MagicMock(return_value=False)
agent.workflow.reset = AsyncMock()
agent.run = AsyncMock(return_value=AgentResponse(messages=[]))
# Constructor inspects WorkflowAgent.workflow internals; bypass setup
# by feeding a configured mock through a normal init.