mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Remove reset
This commit is contained in:
Tao Chen
@@ -1,6 +1,5 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterable, Awaitable
|
||||
from typing import Any, Literal, overload
|
||||
|
||||
@@ -20,7 +19,7 @@ from agent_framework import (
|
||||
WorkflowEvent,
|
||||
WorkflowRunState,
|
||||
)
|
||||
from agent_framework._workflows._agent_executor import AgentExecutorRequest, AgentExecutorResponse
|
||||
from agent_framework._workflows._agent_executor import AgentExecutorResponse
|
||||
from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
|
||||
from agent_framework._workflows._const import GLOBAL_KWARGS_KEY
|
||||
|
||||
@@ -307,108 +306,6 @@ async def test_agent_executor_save_and_restore_state_directly() -> None:
|
||||
assert restored_session.session_id == session.session_id
|
||||
|
||||
|
||||
# region: Tests for AgentExecutor.reset()
|
||||
|
||||
|
||||
async def test_agent_executor_reset_clears_per_run_state() -> None:
|
||||
"""reset() clears cache, conversation snapshot, and pending request/response buffers."""
|
||||
agent = _CountingAgent(id="reset_agent", name="ResetAgent")
|
||||
executor = AgentExecutor(agent, id="reset_exec")
|
||||
|
||||
# Populate every per-run buffer.
|
||||
executor._cache = [Message(role="user", contents=["cached"])] # type: ignore[reportPrivateUsage]
|
||||
executor._full_conversation = [ # type: ignore[reportPrivateUsage]
|
||||
Message(role="user", contents=["prior turn"]),
|
||||
Message(role="assistant", contents=["prior response"]),
|
||||
]
|
||||
pending_request = Content.from_text(text="approve?")
|
||||
executor._pending_agent_requests = {"req-1": pending_request} # type: ignore[reportPrivateUsage]
|
||||
executor._pending_responses_to_agent = [Content.from_text(text="approved")] # type: ignore[reportPrivateUsage]
|
||||
|
||||
await executor.reset()
|
||||
|
||||
assert executor._cache == [] # type: ignore[reportPrivateUsage]
|
||||
assert executor._full_conversation == [] # type: ignore[reportPrivateUsage]
|
||||
assert executor._pending_agent_requests == {} # type: ignore[reportPrivateUsage]
|
||||
assert executor._pending_responses_to_agent == [] # type: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_agent_executor_reset_creates_fresh_session_when_auto_created() -> None:
|
||||
"""reset() replaces the agent session when the executor created it itself."""
|
||||
agent = _CountingAgent(id="reset_session_agent", name="ResetSessionAgent")
|
||||
# No session passed in — executor creates one via agent.create_session().
|
||||
executor = AgentExecutor(agent, id="reset_session_exec")
|
||||
auto_created = executor._session # type: ignore[reportPrivateUsage]
|
||||
auto_created.state["history"] = {"messages": [Message(role="user", contents=["old"])]}
|
||||
|
||||
await executor.reset()
|
||||
|
||||
new_session = executor._session # type: ignore[reportPrivateUsage]
|
||||
assert new_session is not auto_created
|
||||
assert new_session.session_id != auto_created.session_id
|
||||
assert "history" not in new_session.state
|
||||
|
||||
|
||||
async def test_agent_executor_reset_preserves_caller_supplied_session(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""reset() leaves a session passed in via __init__ untouched and warns the caller."""
|
||||
agent = _CountingAgent(id="reset_session_agent", name="ResetSessionAgent")
|
||||
caller_session = AgentSession()
|
||||
history_payload = {"messages": [Message(role="user", contents=["old"])]}
|
||||
caller_session.state["history"] = history_payload
|
||||
executor = AgentExecutor(agent, id="reset_session_exec", session=caller_session)
|
||||
|
||||
assert executor._session is caller_session # type: ignore[reportPrivateUsage]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="agent_framework._workflows._agent_executor"):
|
||||
await executor.reset()
|
||||
|
||||
# Same instance, state untouched — the caller is responsible for managing the session.
|
||||
assert executor._session is caller_session # type: ignore[reportPrivateUsage]
|
||||
assert caller_session.state["history"] is history_payload
|
||||
assert any("Session was supplied by the caller" in record.message for record in caplog.records)
|
||||
|
||||
|
||||
async def test_agent_executor_reset_allows_subsequent_run() -> None:
|
||||
"""After reset(), the executor can be reused for a fresh workflow run without leaking state."""
|
||||
agent = _CountingAgent(id="reset_reuse_agent", name="ResetReuseAgent")
|
||||
executor = AgentExecutor(agent, id="reset_reuse_exec")
|
||||
workflow = WorkflowBuilder(start_executor=executor, output_from=[executor]).build()
|
||||
|
||||
first_outputs: list[WorkflowEvent] = []
|
||||
async for event in workflow.run(
|
||||
AgentExecutorRequest(messages=[Message(role="user", contents=["hello"])]),
|
||||
stream=True,
|
||||
):
|
||||
if event.type == "output":
|
||||
first_outputs.append(event)
|
||||
|
||||
assert first_outputs, "first run should have produced at least one output event"
|
||||
# After a normal run the cache is drained but the conversation snapshot remains.
|
||||
assert executor._cache == [] # type: ignore[reportPrivateUsage]
|
||||
assert executor._full_conversation != [] # type: ignore[reportPrivateUsage]
|
||||
first_session_id = executor._session.session_id # type: ignore[reportPrivateUsage]
|
||||
|
||||
await workflow.reset_for_new_run()
|
||||
|
||||
assert executor._full_conversation == [] # type: ignore[reportPrivateUsage]
|
||||
# Session was auto-created, so reset() rotates it to a fresh one.
|
||||
assert executor._session.session_id != first_session_id # type: ignore[reportPrivateUsage]
|
||||
|
||||
second_outputs: list[WorkflowEvent] = []
|
||||
async for event in workflow.run(
|
||||
AgentExecutorRequest(messages=[Message(role="user", contents=["second"])]),
|
||||
stream=True,
|
||||
):
|
||||
if event.type == "output":
|
||||
second_outputs.append(event)
|
||||
|
||||
assert second_outputs, "second run after reset should have produced at least one output event"
|
||||
assert agent.call_count == 2
|
||||
|
||||
|
||||
# endregion: Tests for AgentExecutor.reset()
|
||||
|
||||
|
||||
async def test_prepare_agent_run_args_extracts_invocation_kwargs() -> None:
|
||||
"""_prepare_agent_run_args extracts function_invocation_kwargs and client_kwargs."""
|
||||
agent = _CountingAgent(id="test_agent", name="TestAgent")
|
||||
|
||||
@@ -1004,53 +1004,3 @@ def test_handler_typevar_error_takes_priority_over_context_error():
|
||||
@handler
|
||||
async def process(self, message: _T, ctx) -> None: # type: ignore[no-untyped-def]
|
||||
pass
|
||||
|
||||
|
||||
# region: Tests for Executor.reset()
|
||||
|
||||
|
||||
async def test_executor_default_reset_is_noop():
|
||||
"""The base Executor.reset() is a no-op and must complete without raising.
|
||||
|
||||
Subclasses that don't carry reset-relevant state should be able to rely on the
|
||||
default implementation.
|
||||
"""
|
||||
|
||||
class StatelessExecutor(Executor):
|
||||
@handler
|
||||
async def handle(self, message: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(message)
|
||||
|
||||
executor_instance = StatelessExecutor(id="stateless")
|
||||
|
||||
# Must complete without raising and return None.
|
||||
assert await executor_instance.reset() is None
|
||||
|
||||
|
||||
async def test_executor_subclass_reset_is_invoked():
|
||||
"""A subclass that overrides reset() can clear its own internal state."""
|
||||
|
||||
class CounterExecutor(Executor):
|
||||
def __init__(self, id: str) -> None:
|
||||
super().__init__(id=id)
|
||||
self.counter = 0
|
||||
self.reset_calls = 0
|
||||
|
||||
@handler
|
||||
async def handle(self, message: int, ctx: WorkflowContext[int]) -> None:
|
||||
self.counter += message
|
||||
|
||||
async def reset(self) -> None:
|
||||
self.counter = 0
|
||||
self.reset_calls += 1
|
||||
|
||||
executor_instance = CounterExecutor(id="counter")
|
||||
executor_instance.counter = 42
|
||||
|
||||
await executor_instance.reset()
|
||||
|
||||
assert executor_instance.counter == 0
|
||||
assert executor_instance.reset_calls == 1
|
||||
|
||||
|
||||
# endregion: Tests for Executor.reset()
|
||||
|
||||
@@ -1111,204 +1111,3 @@ async def test_runner_drains_straggler_events_at_iteration_end():
|
||||
output_events = [e for e in events if e.type == "output"]
|
||||
# We should have output events from both executors
|
||||
assert len(output_events) >= 2
|
||||
|
||||
|
||||
# region: Tests for InProcRunnerContext.reset_for_new_run()
|
||||
|
||||
|
||||
async def test_runner_context_reset_clears_in_flight_messages():
|
||||
"""reset_for_new_run drops queued executor-to-executor messages."""
|
||||
ctx = InProcRunnerContext()
|
||||
await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="src"))
|
||||
|
||||
assert await ctx.has_messages() is True
|
||||
|
||||
ctx.reset_for_new_run()
|
||||
|
||||
assert await ctx.has_messages() is False
|
||||
assert await ctx.drain_messages() == {}
|
||||
|
||||
|
||||
async def test_runner_context_reset_drains_pending_events():
|
||||
"""reset_for_new_run discards any events buffered for streaming."""
|
||||
ctx = InProcRunnerContext()
|
||||
await ctx.add_event(WorkflowEvent.superstep_started(iteration=1))
|
||||
assert await ctx.has_events() is True
|
||||
|
||||
ctx.reset_for_new_run()
|
||||
|
||||
assert await ctx.has_events() is False
|
||||
assert await ctx.drain_events() == []
|
||||
|
||||
|
||||
async def test_runner_context_reset_resets_streaming_flag():
|
||||
"""reset_for_new_run resets streaming back to its non-streaming default."""
|
||||
ctx = InProcRunnerContext()
|
||||
ctx.set_streaming(True)
|
||||
assert ctx.is_streaming() is True
|
||||
|
||||
ctx.reset_for_new_run()
|
||||
|
||||
assert ctx.is_streaming() is False
|
||||
|
||||
|
||||
async def test_runner_context_reset_clears_pending_request_info_events():
|
||||
"""reset_for_new_run clears any pending request_info events tracked for correlation."""
|
||||
ctx = InProcRunnerContext()
|
||||
request_info_event = WorkflowEvent.request_info(
|
||||
request_id="request-123",
|
||||
source_executor_id="source",
|
||||
request_data=MockMessage(data=0),
|
||||
response_type=bool,
|
||||
)
|
||||
await ctx.add_request_info_event(request_info_event)
|
||||
|
||||
assert "request-123" in await ctx.get_pending_request_info_events()
|
||||
|
||||
ctx.reset_for_new_run()
|
||||
|
||||
assert await ctx.get_pending_request_info_events() == {}
|
||||
|
||||
|
||||
# endregion: Tests for InProcRunnerContext.reset_for_new_run()
|
||||
|
||||
|
||||
# region: Tests for Runner.reset_for_new_run()
|
||||
|
||||
|
||||
async def test_runner_reset_for_new_run_resets_iteration_count():
|
||||
"""reset_for_new_run resets the iteration counter back to zero."""
|
||||
runner = _make_runner()
|
||||
runner._iteration = 7 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
await runner.reset_for_new_run()
|
||||
|
||||
assert runner._iteration == 0 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_runner_reset_for_new_run_clears_shared_state():
|
||||
"""reset_for_new_run wipes both committed and pending entries from shared state."""
|
||||
state = State()
|
||||
state.set("committed_key", "committed_value")
|
||||
state.commit()
|
||||
state.set("pending_key", "pending_value") # uncommitted
|
||||
|
||||
runner = Runner(
|
||||
[],
|
||||
{},
|
||||
state,
|
||||
InProcRunnerContext(),
|
||||
"test_name",
|
||||
graph_signature_hash="test_hash",
|
||||
)
|
||||
|
||||
await runner.reset_for_new_run()
|
||||
|
||||
assert state.get("committed_key") is None
|
||||
assert state.get("pending_key") is None
|
||||
assert state.has("committed_key") is False
|
||||
assert state.has("pending_key") is False
|
||||
|
||||
|
||||
async def test_runner_reset_for_new_run_clears_resumed_from_checkpoint_flag():
|
||||
"""reset_for_new_run clears the flag set by restore_from_checkpoint."""
|
||||
runner = _make_runner()
|
||||
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]
|
||||
assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
await runner.reset_for_new_run()
|
||||
|
||||
assert runner._resumed_from_checkpoint is False # pyright: ignore[reportPrivateUsage]
|
||||
# And the iteration count restored from the checkpoint must be wiped, too.
|
||||
assert runner._iteration == 0 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_runner_reset_for_new_run_invokes_executor_reset_for_each_executor():
|
||||
"""reset_for_new_run calls reset() on every registered executor exactly once."""
|
||||
|
||||
class TrackingExecutor(MockExecutor):
|
||||
def __init__(self, id: str) -> None:
|
||||
super().__init__(id=id)
|
||||
self.reset_calls = 0
|
||||
|
||||
async def reset(self) -> None:
|
||||
self.reset_calls += 1
|
||||
|
||||
executor_a = TrackingExecutor(id="executor_a")
|
||||
executor_b = TrackingExecutor(id="executor_b")
|
||||
|
||||
runner = Runner(
|
||||
[],
|
||||
{executor_a.id: executor_a, executor_b.id: executor_b},
|
||||
State(),
|
||||
InProcRunnerContext(),
|
||||
"test_name",
|
||||
graph_signature_hash="test_hash",
|
||||
)
|
||||
|
||||
await runner.reset_for_new_run()
|
||||
|
||||
assert executor_a.reset_calls == 1
|
||||
assert executor_b.reset_calls == 1
|
||||
|
||||
|
||||
async def test_runner_reset_for_new_run_resets_runner_context():
|
||||
"""reset_for_new_run forwards the reset to the underlying runner context."""
|
||||
ctx = InProcRunnerContext()
|
||||
await ctx.send_message(WorkflowMessage(data=MockMessage(data=0), source_id="src"))
|
||||
await ctx.add_event(WorkflowEvent.superstep_started(iteration=1))
|
||||
ctx.set_streaming(True)
|
||||
|
||||
runner = Runner([], {}, State(), ctx, "test_name", graph_signature_hash="test_hash")
|
||||
|
||||
await runner.reset_for_new_run()
|
||||
|
||||
assert await ctx.has_messages() is False
|
||||
assert await ctx.has_events() is False
|
||||
assert ctx.is_streaming() is False
|
||||
|
||||
|
||||
async def test_runner_can_run_again_after_reset_for_new_run():
|
||||
"""After reset_for_new_run the runner can be reserved and converge a fresh workload."""
|
||||
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()
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
|
||||
|
||||
# First run: drives MockExecutor's loop until it yields the terminal value.
|
||||
await executor_a.execute(MockMessage(data=0), ["START"], state, ctx)
|
||||
async for _ in runner.run_until_convergence():
|
||||
pass
|
||||
assert runner._iteration == 10 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
await runner.reset_for_new_run()
|
||||
|
||||
# Second run: must succeed cleanly using the same runner instance.
|
||||
await executor_a.execute(MockMessage(data=0), ["START"], state, ctx)
|
||||
second_run_outputs: list[int] = []
|
||||
async for event in runner.run_until_convergence():
|
||||
if event.type == "output":
|
||||
second_run_outputs.append(event.data)
|
||||
|
||||
assert second_run_outputs == [10]
|
||||
assert runner._iteration == 10 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
# endregion: Tests for Runner.reset_for_new_run()
|
||||
|
||||
@@ -689,89 +689,3 @@ async def test_sub_workflow_intermediate_outputs_propagate_to_parent() -> None:
|
||||
|
||||
# The parent's own terminal output is unaffected.
|
||||
assert any(e.executor_id == "parent_sink" and e.data == "final: hello" for e in output_events)
|
||||
|
||||
|
||||
# region: Tests for WorkflowExecutor.reset()
|
||||
|
||||
|
||||
async def test_workflow_executor_reset_clears_execution_state() -> None:
|
||||
"""reset() clears the WorkflowExecutor's per-run execution contexts and request mappings."""
|
||||
validation_workflow = create_email_validation_workflow()
|
||||
parent = Coordinator()
|
||||
workflow_executor = WorkflowExecutor(validation_workflow, "email_validation_workflow")
|
||||
|
||||
main_workflow = (
|
||||
WorkflowBuilder(start_executor=parent)
|
||||
.add_edge(parent, workflow_executor)
|
||||
.add_edge(workflow_executor, parent)
|
||||
.build()
|
||||
)
|
||||
|
||||
# First run pauses with a pending request from the sub-workflow.
|
||||
result = await main_workflow.run("test@example.com")
|
||||
assert len(result.get_request_info_events()) == 1
|
||||
assert len(workflow_executor._execution_contexts) == 1 # type: ignore[reportPrivateUsage]
|
||||
assert len(workflow_executor._request_to_execution) == 1 # type: ignore[reportPrivateUsage]
|
||||
|
||||
await main_workflow.reset_for_new_run()
|
||||
|
||||
assert workflow_executor._execution_contexts == {} # type: ignore[reportPrivateUsage]
|
||||
assert workflow_executor._request_to_execution == {} # type: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_workflow_executor_reset_resets_wrapped_workflow() -> None:
|
||||
"""reset() recursively resets the wrapped workflow (runner iteration counter cleared)."""
|
||||
validation_workflow = create_email_validation_workflow()
|
||||
parent = Coordinator()
|
||||
workflow_executor = WorkflowExecutor(validation_workflow, "email_validation_workflow")
|
||||
|
||||
main_workflow = (
|
||||
WorkflowBuilder(start_executor=parent)
|
||||
.add_edge(parent, workflow_executor)
|
||||
.add_edge(workflow_executor, parent)
|
||||
.build()
|
||||
)
|
||||
|
||||
await main_workflow.run("test@example.com")
|
||||
# The sub-workflow's runner advanced past iteration 0 during execution.
|
||||
assert validation_workflow._runner._iteration > 0 # type: ignore[reportPrivateUsage]
|
||||
|
||||
await main_workflow.reset_for_new_run()
|
||||
|
||||
# The wrapped workflow's runner was reset along with the parent.
|
||||
assert validation_workflow._runner._iteration == 0 # type: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_workflow_executor_reset_allows_subsequent_run() -> None:
|
||||
"""After reset(), the parent + WorkflowExecutor can be reused for a fresh run with no leakage."""
|
||||
validation_workflow = create_email_validation_workflow()
|
||||
parent = Coordinator()
|
||||
workflow_executor = WorkflowExecutor(validation_workflow, "email_validation_workflow")
|
||||
|
||||
main_workflow = (
|
||||
WorkflowBuilder(start_executor=parent)
|
||||
.add_edge(parent, workflow_executor)
|
||||
.add_edge(workflow_executor, parent)
|
||||
.build()
|
||||
)
|
||||
|
||||
first_result = await main_workflow.run("first@example.com")
|
||||
assert len(first_result.get_request_info_events()) == 1
|
||||
|
||||
await main_workflow.reset_for_new_run()
|
||||
|
||||
# State on the WorkflowExecutor and parent's pending-request bookkeeping is clean.
|
||||
assert workflow_executor._execution_contexts == {} # type: ignore[reportPrivateUsage]
|
||||
assert workflow_executor._request_to_execution == {} # type: ignore[reportPrivateUsage]
|
||||
|
||||
second_result = await main_workflow.run("second@example.com")
|
||||
second_requests = second_result.get_request_info_events()
|
||||
assert len(second_requests) == 1
|
||||
assert isinstance(second_requests[0].data, DomainCheckRequest)
|
||||
# Confirm the new run produced a request from the second email, not the cached first one.
|
||||
assert second_requests[0].data.email == "second@example.com"
|
||||
# And the WorkflowExecutor is now tracking exactly one fresh execution.
|
||||
assert len(workflow_executor._execution_contexts) == 1 # type: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
# endregion: Tests for WorkflowExecutor.reset()
|
||||
|
||||
@@ -29,7 +29,6 @@ from agent_framework import (
|
||||
WorkflowEvent,
|
||||
WorkflowException,
|
||||
WorkflowMessage,
|
||||
WorkflowRunnerException,
|
||||
WorkflowRunState,
|
||||
handler,
|
||||
response_handler,
|
||||
@@ -1354,144 +1353,3 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region: Tests for Workflow.reset_for_new_run()
|
||||
|
||||
|
||||
async def test_workflow_reset_for_new_run_allows_subsequent_run() -> None:
|
||||
"""After reset_for_new_run() the same workflow instance can be run again from scratch."""
|
||||
executor_a = IncrementExecutor(id="executor_a")
|
||||
executor_b = IncrementExecutor(id="executor_b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=executor_a, output_from=[executor_a, executor_b])
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.build()
|
||||
)
|
||||
|
||||
first = await workflow.run(NumberMessage(data=0))
|
||||
assert first.get_outputs() == [10]
|
||||
|
||||
await workflow.reset_for_new_run()
|
||||
|
||||
second = await workflow.run(NumberMessage(data=0))
|
||||
assert second.get_outputs() == [10]
|
||||
|
||||
|
||||
async def test_workflow_reset_for_new_run_clears_workflow_state() -> None:
|
||||
"""reset_for_new_run() clears values that executors persisted in shared workflow state."""
|
||||
|
||||
class StateWritingExecutor(Executor):
|
||||
@handler
|
||||
async def handle(self, message: NumberMessage, ctx: WorkflowContext[Any, int]) -> None:
|
||||
previous = ctx.get_state("seen") or 0
|
||||
ctx.set_state("seen", previous + 1)
|
||||
await ctx.yield_output(previous + 1)
|
||||
|
||||
state_writer = StateWritingExecutor(id="state_writer")
|
||||
workflow = WorkflowBuilder(start_executor=state_writer, output_from=[state_writer]).build()
|
||||
|
||||
first = await workflow.run(NumberMessage(data=1))
|
||||
assert first.get_outputs() == [1]
|
||||
# State was persisted by the executor.
|
||||
assert workflow._runner.state.get("seen") == 1 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
await workflow.reset_for_new_run()
|
||||
|
||||
# The runner's shared state has been wiped.
|
||||
assert workflow._runner.state.get("seen") is None # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
second = await workflow.run(NumberMessage(data=1))
|
||||
# Counter started fresh from 0 again; output is 1, not 2.
|
||||
assert second.get_outputs() == [1]
|
||||
|
||||
|
||||
async def test_workflow_reset_for_new_run_invokes_executor_reset_hook() -> None:
|
||||
"""reset_for_new_run() calls Executor.reset() on every executor in the workflow."""
|
||||
|
||||
class ResettableExecutor(Executor):
|
||||
def __init__(self, id: str) -> None:
|
||||
super().__init__(id=id)
|
||||
self.reset_calls = 0
|
||||
self.handled = 0
|
||||
|
||||
@handler
|
||||
async def handle(self, message: NumberMessage, ctx: WorkflowContext[Any, int]) -> None:
|
||||
self.handled += 1
|
||||
await ctx.yield_output(self.handled)
|
||||
|
||||
async def reset(self) -> None:
|
||||
self.reset_calls += 1
|
||||
self.handled = 0
|
||||
|
||||
executor = ResettableExecutor(id="resettable")
|
||||
workflow = WorkflowBuilder(start_executor=executor, output_from=[executor]).build()
|
||||
|
||||
await workflow.run(NumberMessage(data=1))
|
||||
assert executor.handled == 1
|
||||
assert executor.reset_calls == 0
|
||||
|
||||
await workflow.reset_for_new_run()
|
||||
|
||||
assert executor.reset_calls == 1
|
||||
# The executor's own counter was wiped by its overridden reset().
|
||||
assert executor.handled == 0
|
||||
|
||||
|
||||
async def test_workflow_reset_for_new_run_resets_runner_iteration_counter() -> None:
|
||||
"""reset_for_new_run() drops the iteration counter accumulated during a prior run."""
|
||||
executor_a = IncrementExecutor(id="executor_a")
|
||||
executor_b = IncrementExecutor(id="executor_b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=executor_a, output_from=[executor_a, executor_b])
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.build()
|
||||
)
|
||||
|
||||
await workflow.run(NumberMessage(data=0))
|
||||
assert workflow._runner._iteration > 0 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
await workflow.reset_for_new_run()
|
||||
|
||||
assert workflow._runner._iteration == 0 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_workflow_reset_for_new_run_rejected_during_streaming_run() -> None:
|
||||
"""reset_for_new_run() raises WorkflowRunnerException while a streaming run is in progress."""
|
||||
executor_a = IncrementExecutor(id="executor_a")
|
||||
executor_b = IncrementExecutor(id="executor_b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=executor_a, output_from=[executor_a, executor_b])
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.build()
|
||||
)
|
||||
|
||||
async def consume_stream_slowly() -> list[WorkflowEvent]:
|
||||
events: list[WorkflowEvent] = []
|
||||
async for event in workflow.run(NumberMessage(data=0), stream=True):
|
||||
events.append(event)
|
||||
await asyncio.sleep(0.01)
|
||||
return events
|
||||
|
||||
task = asyncio.create_task(consume_stream_slowly())
|
||||
# Let the streaming run start.
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
try:
|
||||
with pytest.raises(WorkflowRunnerException, match="Cannot reset the workflow while a run is in progress"):
|
||||
await workflow.reset_for_new_run()
|
||||
finally:
|
||||
await task
|
||||
|
||||
# After the run completes, reset succeeds again.
|
||||
await workflow.reset_for_new_run()
|
||||
assert workflow._runner._iteration == 0 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
Reference in New Issue
Block a user