Python: [Breaking] Remove WorkflowCompletedEvent, introduce workflow output and migrate to ctx.yield_output() + a huge refactoring (#845)

* Introduce input and output types for executor and workflow

* WorkflowOutputContext handles two types

* Remove can_handle_types from Executor

* Update validation

* Move workflow executor

* Move workflow executor

* Fix issues in WorkflowExecutor

* refactor executor

* update execute signature to create workflow context within Executor

* fix simple sub workflow test; fix validation

* fix output types in WorkflowExecutor

* fix issue in Executor handling of SubWorkflowRequestInfo

* update tests to use proper workflow output

* update orchestration patterns to use output

* Update sample -- not finished

* Update python/packages/main/tests/workflow/test_workflow_states.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update python/packages/main/tests/workflow/test_concurrent.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* address comments

* WorkflowOutputContext --> WorkflowContext

* remove WorkflowCompletedEvent

* update samples

* Update doc string for important classes; update WorkflowExecutor to support concurrent execution

* use Never instead of None for default type

* Update usage of WorkflowContext[None to WorkflowContext[Never

* address comments

* remove filter for None

* address comments, minor fixes

* quality of life improvement on interceptor types

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Eric Zhu
2025-09-23 13:52:53 -07:00
committed by GitHub
Unverified
parent 0f913bcdeb
commit 2133043f11
67 changed files with 2564 additions and 1648 deletions
@@ -1,8 +1,9 @@
# Copyright (c) Microsoft. All rights reserved.
import pytest
from typing_extensions import Never
from agent_framework import WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler
from agent_framework import WorkflowBuilder, WorkflowContext, WorkflowRunState, WorkflowStatusEvent, handler
from agent_framework._workflow._checkpoint import InMemoryCheckpointStorage
from agent_framework._workflow._executor import Executor
@@ -15,8 +16,8 @@ class StartExecutor(Executor):
class FinishExecutor(Executor):
@handler
async def finish(self, message: str, ctx: WorkflowContext[None]) -> None:
await ctx.add_event(WorkflowCompletedEvent(message))
async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None:
await ctx.yield_output(message)
def build_workflow(storage: InMemoryCheckpointStorage, finish_id: str = "finish"):
@@ -70,4 +71,4 @@ async def test_resume_succeeds_when_graph_matches() -> None:
)
]
assert any(isinstance(event, WorkflowCompletedEvent) for event in events)
assert any(isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE for event in events)
@@ -12,8 +12,10 @@ from agent_framework import (
ConcurrentBuilder,
Executor,
Role,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
@@ -57,15 +59,19 @@ async def test_concurrent_default_aggregator_emits_single_user_and_assistants()
wf = ConcurrentBuilder().participants([e1, e2, e3]).build()
completed: WorkflowCompletedEvent | None = None
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run_stream("prompt: hello world"):
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
output = cast(list[ChatMessage], ev.data)
if completed and output is not None:
break
assert completed is not None
assert isinstance(completed.data, list)
messages: list[ChatMessage] = cast(list[ChatMessage], completed.data) # type: ignore
assert completed
assert output is not None
messages: list[ChatMessage] = output
# Expect one user message + one assistant message per participant
assert len(messages) == 1 + 3
@@ -91,16 +97,21 @@ async def test_concurrent_custom_aggregator_callback_is_used() -> None:
wf = ConcurrentBuilder().participants([e1, e2]).with_aggregator(summarize).build()
completed: WorkflowCompletedEvent | None = None
completed = False
output: str | None = None
async for ev in wf.run_stream("prompt: custom"):
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
output = cast(str, ev.data)
if completed and output is not None:
break
assert completed is not None
assert completed
assert output is not None
# Custom aggregator returns a string payload
assert isinstance(completed.data, str)
assert completed.data == "One | Two"
assert isinstance(output, str)
assert output == "One | Two"
async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None:
@@ -108,7 +119,7 @@ async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None:
e2 = _FakeAgentExec("agentB", "Two")
# Sync callback with ctx parameter (should run via asyncio.to_thread)
def summarize_sync(results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> str: # type: ignore[unused-argument]
def summarize_sync(results: list[AgentExecutorResponse], _ctx: WorkflowContext[Any]) -> str: # type: ignore[unused-argument]
texts: list[str] = []
for r in results:
msgs: list[ChatMessage] = r.agent_run_response.messages
@@ -117,15 +128,20 @@ async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None:
wf = ConcurrentBuilder().participants([e1, e2]).with_aggregator(summarize_sync).build()
completed: WorkflowCompletedEvent | None = None
completed = False
output: str | None = None
async for ev in wf.run_stream("prompt: custom sync"):
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
output = cast(str, ev.data)
if completed and output is not None:
break
assert completed is not None
assert isinstance(completed.data, str)
assert completed.data == "One | Two"
assert completed
assert output is not None
assert isinstance(output, str)
assert output == "One | Two"
def test_concurrent_custom_aggregator_uses_callback_name_for_id() -> None:
@@ -4,6 +4,7 @@ from collections.abc import AsyncIterable
from typing import Any
from pydantic import PrivateAttr
from typing_extensions import Never
from agent_framework import (
AgentExecutor,
@@ -16,11 +17,13 @@ from agent_framework import (
SequentialBuilder,
TextContent,
WorkflowBuilder,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
from agent_framework._workflow._executor import AgentExecutorResponse, Executor
from agent_framework._workflow._workflow_context import WorkflowContext
class _SimpleAgent(BaseAgent):
@@ -54,19 +57,17 @@ class _CaptureFullConversation(Executor):
"""Captures AgentExecutorResponse.full_conversation and completes the workflow."""
@handler
async def capture(self, response: AgentExecutorResponse, ctx: WorkflowContext[None]) -> None:
async def capture(self, response: AgentExecutorResponse, ctx: WorkflowContext[Never, dict]) -> None:
full = response.full_conversation
# The AgentExecutor contract guarantees full_conversation is populated.
assert full is not None
await ctx.add_event(
WorkflowCompletedEvent(
data={
"length": len(full),
"roles": [m.role for m in full],
"texts": [m.text for m in full],
}
)
)
payload = {
"length": len(full),
"roles": [m.role for m in full],
"texts": [m.text for m in full],
}
await ctx.yield_output(payload)
pass
async def test_agent_executor_populates_full_conversation_non_streaming() -> None:
@@ -78,15 +79,20 @@ async def test_agent_executor_populates_full_conversation_non_streaming() -> Non
wf = WorkflowBuilder().set_start_executor(agent_exec).add_edge(agent_exec, capturer).build()
# Act: run with a simple user prompt
completed: WorkflowCompletedEvent | None = None
completed = False
output: dict | None = None
async for ev in wf.run_stream("hello world"):
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
output = ev.data # type: ignore[assignment]
if completed and output is not None:
break
# Assert: full_conversation contains [user("hello world"), assistant("agent-reply")]
assert completed is not None
payload = completed.data # type: ignore[assignment]
assert completed
assert output is not None
payload = output
assert isinstance(payload, dict)
assert payload["length"] == 2
assert payload["roles"][0] == Role.USER and "hello world" in (payload["texts"][0] or "")
@@ -148,7 +154,7 @@ async def test_sequential_adapter_uses_full_conversation() -> None:
# Act
async for ev in wf.run_stream("hello seq"):
if isinstance(ev, WorkflowCompletedEvent):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
break
# Assert: second agent should have seen the user prompt and A1's assistant reply
@@ -3,11 +3,11 @@
from typing import Any
import pytest
from typing_extensions import Never
from agent_framework import (
FunctionExecutor,
WorkflowBuilder,
WorkflowCompletedEvent,
WorkflowContext,
executor,
)
@@ -28,9 +28,9 @@ class TestFunctionExecutor:
assert len(func_exec._handlers) == 1
assert str in func_exec._handlers
# Check instance handler spec was created
assert len(func_exec._instance_handler_specs) == 1
spec = func_exec._instance_handler_specs[0]
# Check handler spec was created
assert len(func_exec._handler_specs) == 1
spec = func_exec._handler_specs[0]
assert spec["name"] == "process_string"
assert spec["message_type"] is str
assert spec["output_types"] == [str]
@@ -47,7 +47,7 @@ class TestFunctionExecutor:
assert int in process_int._handlers
# Check spec
spec = process_int._instance_handler_specs[0]
spec = process_int._handler_specs[0]
assert spec["message_type"] is int
assert spec["output_types"] == [int]
@@ -81,7 +81,7 @@ class TestFunctionExecutor:
assert int in simple_no_parens._handlers
def test_union_output_types(self):
"""Test that union output types are properly inferred."""
"""Test that union output types are properly inferred for both messages and workflow outputs."""
@executor
async def multi_output(text: str, ctx: WorkflowContext[str | int]) -> None:
@@ -90,29 +90,56 @@ class TestFunctionExecutor:
else:
await ctx.send_message(text.upper())
spec = multi_output._instance_handler_specs[0]
spec = multi_output._handler_specs[0]
assert set(spec["output_types"]) == {str, int}
assert spec["workflow_output_types"] == [] # No workflow outputs defined
# Test union types for workflow outputs too
@executor
async def multi_workflow_output(data: str, ctx: WorkflowContext[Never, str | int | bool]) -> None:
if data.isdigit():
await ctx.yield_output(int(data))
elif data.lower() in ("true", "false"):
await ctx.yield_output(data.lower() == "true")
else:
await ctx.yield_output(data.upper())
workflow_spec = multi_workflow_output._handler_specs[0]
assert workflow_spec["output_types"] == [] # None means no message outputs
assert set(workflow_spec["workflow_output_types"]) == {str, int, bool}
def test_none_output_type(self):
"""Test WorkflowContext[None] produces empty output types."""
"""Test WorkflowContext produces empty output types."""
@executor
async def no_output(data: Any, ctx: WorkflowContext[None]) -> None:
async def no_output(data: Any, ctx: WorkflowContext) -> None:
# This executor doesn't send any messages
pass
spec = no_output._instance_handler_specs[0]
spec = no_output._handler_specs[0]
assert spec["output_types"] == []
assert spec["workflow_output_types"] == [] # No workflow outputs defined
def test_any_output_type(self):
"""Test WorkflowContext[Any] produces empty output types."""
"""Test WorkflowContext[Any] and WorkflowContext[Any, Any] produce Any output types."""
@executor
async def any_output(data: str, ctx: WorkflowContext[Any]) -> None:
await ctx.send_message("result")
spec = any_output._instance_handler_specs[0]
assert spec["output_types"] == []
spec = any_output._handler_specs[0]
assert spec["output_types"] == [Any]
assert spec["workflow_output_types"] == [] # No workflow outputs defined
# Test both parameters as Any
@executor
async def any_both_output(data: str, ctx: WorkflowContext[Any, Any]) -> None:
await ctx.send_message("message")
await ctx.yield_output("workflow_output")
both_spec = any_both_output._handler_specs[0]
assert both_spec["output_types"] == [Any]
assert both_spec["workflow_output_types"] == [Any]
def test_validation_errors(self):
"""Test various validation errors in function signatures."""
@@ -121,13 +148,17 @@ class TestFunctionExecutor:
async def no_params() -> None:
pass
with pytest.raises(ValueError, match="one or two parameters"):
with pytest.raises(
ValueError, match="must have \\(message: T\\) or \\(message: T, ctx: WorkflowContext\\[U\\]\\)"
):
FunctionExecutor(no_params) # type: ignore
async def too_many_params(data: str, ctx: WorkflowContext[str], extra: int) -> None:
pass
with pytest.raises(ValueError, match="one or two parameters"):
with pytest.raises(
ValueError, match="must have \\(message: T\\) or \\(message: T, ctx: WorkflowContext\\[U\\]\\)"
):
FunctionExecutor(too_many_params) # type: ignore
# Missing message type annotation
@@ -141,22 +172,24 @@ class TestFunctionExecutor:
async def no_ctx_type(data: str, ctx) -> None: # type: ignore
pass
with pytest.raises(ValueError, match="annotated as WorkflowContext"):
with pytest.raises(ValueError, match="must have a WorkflowContext"):
FunctionExecutor(no_ctx_type) # type: ignore
# Wrong ctx type
async def wrong_ctx_type(data: str, ctx: str) -> None: # type: ignore
pass
with pytest.raises(ValueError, match="WorkflowContext\\[T\\]"):
with pytest.raises(ValueError, match="must be annotated as WorkflowContext"):
FunctionExecutor(wrong_ctx_type) # type: ignore
# Unparameterized WorkflowContext
# Unparameterized WorkflowContext is now allowed
async def unparameterized_ctx(data: str, ctx: WorkflowContext) -> None: # type: ignore
pass
with pytest.raises(ValueError, match="concrete T"):
FunctionExecutor(unparameterized_ctx) # type: ignore
# This should now succeed since unparameterized WorkflowContext is allowed
executor = FunctionExecutor(unparameterized_ctx)
assert executor.output_types == [] # Unparameterized has no inferred types
assert executor.workflow_output_types == [] # No workflow output types
async def test_execution_in_workflow(self):
"""Test that FunctionExecutor works properly in a workflow."""
@@ -167,18 +200,28 @@ class TestFunctionExecutor:
await ctx.send_message(result)
@executor(id="reverse")
async def reverse_text(text: str, ctx: WorkflowContext[Any]) -> None:
async def reverse_text(text: str, ctx: WorkflowContext[Any, str]) -> None:
result = text[::-1]
await ctx.add_event(WorkflowCompletedEvent(result))
await ctx.yield_output(result)
# Verify type inference for both executors
upper_spec = to_upper._handler_specs[0]
assert upper_spec["output_types"] == [str]
assert upper_spec["workflow_output_types"] == [] # No workflow outputs
reverse_spec = reverse_text._handler_specs[0]
assert reverse_spec["output_types"] == [Any] # First parameter is Any
assert reverse_spec["workflow_output_types"] == [str] # Second parameter is str
workflow = WorkflowBuilder().add_edge(to_upper, reverse_text).set_start_executor(to_upper).build()
# Run workflow
events = await workflow.run("hello world")
completed = events.get_completed_event()
outputs = events.get_outputs()
assert completed is not None
assert completed.data == "DLROW OLLEH"
# Assert that we got the expected output
assert len(outputs) == 1
assert outputs[0] == "DLROW OLLEH"
def test_can_handle_method(self):
"""Test that can_handle method works with instance handlers."""
@@ -204,12 +247,13 @@ class TestFunctionExecutor:
await ctx.send_message(message)
with pytest.raises(ValueError, match="Handler for type .* already registered"):
func_exec.register_instance_handler(
func_exec._register_instance_handler(
name="second",
func=second_handler,
message_type=str,
ctx_annotation=WorkflowContext[str],
output_types=[str],
workflow_output_types=[],
)
def test_complex_type_annotations(self):
@@ -220,7 +264,7 @@ class TestFunctionExecutor:
result = {item: len(item) for item in items}
await ctx.send_message(result)
spec = process_list._instance_handler_specs[0]
spec = process_list._handler_specs[0]
assert spec["message_type"] == list[str]
assert spec["output_types"] == [dict[str, int]]
@@ -236,7 +280,7 @@ class TestFunctionExecutor:
assert str in process_simple._handlers
# Check spec - single parameter functions have no output types since they can't send messages
spec = process_simple._instance_handler_specs[0]
spec = process_simple._handler_specs[0]
assert spec["message_type"] is str
assert spec["output_types"] == []
assert spec["ctx_annotation"] is None
@@ -296,7 +340,7 @@ class TestFunctionExecutor:
assert str in process_sync._handlers
# Check spec - sync single parameter functions have no output types
spec = process_sync._instance_handler_specs[0]
spec = process_sync._handler_specs[0]
assert spec["message_type"] is str
assert spec["output_types"] == []
assert spec["ctx_annotation"] is None
@@ -314,7 +358,7 @@ class TestFunctionExecutor:
assert int in sync_with_ctx._handlers
# Check spec - sync functions with context can infer output types
spec = sync_with_ctx._instance_handler_specs[0]
spec = sync_with_ctx._handler_specs[0]
assert spec["message_type"] is int
assert spec["output_types"] == [int]
@@ -385,9 +429,18 @@ class TestFunctionExecutor:
# In practice, the wrapper handles the async conversion
@executor(id="async_reverse")
async def reverse_async(text: str, ctx: WorkflowContext[Any]):
async def reverse_async(text: str, ctx: WorkflowContext[Any, str]):
result = text[::-1]
await ctx.add_event(WorkflowCompletedEvent(result))
await ctx.yield_output(result)
# Verify type inference for sync and async functions
sync_spec = to_upper_sync._handler_specs[0]
assert sync_spec["output_types"] == [str]
assert sync_spec["workflow_output_types"] == [] # No workflow outputs
async_spec = reverse_async._handler_specs[0]
assert async_spec["output_types"] == [Any] # First parameter is Any
assert async_spec["workflow_output_types"] == [str] # Second parameter is str
# Verify the executors can handle their input types
assert to_upper_sync.can_handle("hello")
@@ -23,9 +23,11 @@ from agent_framework import (
RequestInfoEvent,
Role,
TextContent,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowEvent, # type: ignore # noqa: E402
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
from agent_framework._agents import BaseAgent
@@ -169,15 +171,20 @@ async def test_magentic_workflow_plan_review_approval_to_completion():
req_event = ev
assert req_event is not None
completed: WorkflowCompletedEvent | None = None
completed = False
output: ChatMessage | None = None
async for ev in wf.send_responses_streaming({
req_event.request_id: MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE)
}):
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
output = ev.data # type: ignore[assignment]
if completed and output is not None:
break
assert completed is not None
assert isinstance(getattr(completed, "data", None), ChatMessage)
assert completed
assert output is not None
assert isinstance(output, ChatMessage)
async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds():
@@ -210,7 +217,7 @@ async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds()
# Reply APPROVE with comments (no edited text). Expect one replan and no second review round.
saw_second_review = False
completed: WorkflowCompletedEvent | None = None
completed = False
async for ev in wf.send_responses_streaming({
req_event.request_id: MagenticPlanReviewReply(
decision=MagenticPlanReviewDecision.APPROVE,
@@ -219,11 +226,11 @@ async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds()
}):
if isinstance(ev, RequestInfoEvent) and ev.request_type is MagenticPlanReviewRequest:
saw_second_review = True
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
break
assert completed is not None
assert completed
assert manager.replan_count >= 1
assert saw_second_review is False
# Replan from FakeManager updates facts/plan to include A2 / Do Z
@@ -245,9 +252,14 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result():
if len(events) > 50:
break
completed = next((e for e in events if isinstance(e, WorkflowCompletedEvent)), None)
assert completed is not None
data = getattr(completed, "data", None)
idle_status = next(
(e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE), None
)
assert idle_status is not None
# Check that we got workflow output via WorkflowOutputEvent
output_event = next((e for e in events if isinstance(e, WorkflowOutputEvent)), None)
assert output_event is not None
data = output_event.data
assert isinstance(data, ChatMessage)
assert data.role == Role.ASSISTANT
@@ -9,10 +9,11 @@ from agent_framework import (
AgentExecutorResponse,
AgentRunResponse,
Executor,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowEvent,
WorkflowEventSource,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
from agent_framework._workflow._edge import SingleEdgeGroup
@@ -32,11 +33,12 @@ class MockExecutor(Executor):
"""A mock executor for testing purposes."""
@handler
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None:
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
if message.data < 10:
await ctx.send_message(MockMessage(data=message.data + 1))
else:
await ctx.add_event(WorkflowCompletedEvent(data=message.data))
await ctx.yield_output(message.data)
pass
def test_create_runner():
@@ -77,18 +79,14 @@ async def test_runner_run_until_convergence():
result: int | None = None
await executor_a.execute(
MockMessage(data=0),
WorkflowContext(
executor_id=executor_a.id,
source_executor_ids=["START"],
shared_state=shared_state,
runner_context=ctx,
),
["START"], # source_executor_ids
shared_state, # shared_state
ctx, # runner_context
)
async for event in runner.run_until_convergence():
assert isinstance(event, WorkflowEvent)
if isinstance(event, WorkflowCompletedEvent):
if isinstance(event, WorkflowOutputEvent):
result = event.data
assert event.origin is WorkflowEventSource.EXECUTOR
assert result is not None and result == 10
@@ -112,16 +110,13 @@ async def test_runner_run_until_convergence_not_completed():
await executor_a.execute(
MockMessage(data=0),
WorkflowContext(
executor_id=executor_a.id,
source_executor_ids=["START"],
shared_state=shared_state,
runner_context=ctx,
),
["START"], # source_executor_ids
shared_state, # shared_state
ctx, # runner_context
)
with pytest.raises(RuntimeError, match="Runner did not converge after 5 iterations."):
async for event in runner.run_until_convergence():
assert not isinstance(event, WorkflowCompletedEvent)
assert not isinstance(event, WorkflowStatusEvent) or event.state != WorkflowRunState.IDLE
async def test_runner_already_running():
@@ -143,12 +138,9 @@ async def test_runner_already_running():
await executor_a.execute(
MockMessage(data=0),
WorkflowContext(
executor_id=executor_a.id,
source_executor_ids=["START"],
shared_state=shared_state,
runner_context=ctx,
),
["START"], # source_executor_ids
shared_state, # shared_state
ctx, # runner_context
)
with pytest.raises(RuntimeError, match="Runner is already running."):
@@ -172,6 +164,6 @@ async def test_runner_emits_runner_completion_for_agent_response_without_targets
)
events: list[WorkflowEvent] = [event async for event in runner.run_until_convergence()]
completions = [e for e in events if isinstance(e, WorkflowCompletedEvent)]
assert completions
assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in completions)
# The runner should complete without errors when handling AgentExecutorResponse without targets
# No specific events are expected since there are no executors to process the message
assert isinstance(events, list) # Just verify the runner completed without errors
@@ -15,8 +15,10 @@ from agent_framework import (
Role,
SequentialBuilder,
TextContent,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
@@ -66,15 +68,20 @@ async def test_sequential_agents_append_to_context() -> None:
wf = SequentialBuilder().participants([a1, a2]).build()
completed: WorkflowCompletedEvent | None = None
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run_stream("hello sequential"):
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
output = ev.data # type: ignore[assignment]
if completed and output is not None:
break
assert completed is not None
assert isinstance(completed.data, list)
msgs: list[ChatMessage] = completed.data # type: ignore[assignment]
assert completed
assert output is not None
assert isinstance(output, list)
msgs: list[ChatMessage] = output
assert len(msgs) == 3
assert msgs[0].role == Role.USER and "hello sequential" in msgs[0].text
assert msgs[1].role == Role.ASSISTANT and (msgs[1].author_name == "A1" or True)
@@ -89,14 +96,19 @@ async def test_sequential_with_custom_executor_summary() -> None:
wf = SequentialBuilder().participants([a1, summarizer]).build()
completed: WorkflowCompletedEvent | None = None
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run_stream("topic X"):
if isinstance(ev, WorkflowCompletedEvent):
completed = ev
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
output = ev.data # type: ignore[assignment]
if completed and output is not None:
break
assert completed is not None
msgs: list[ChatMessage] = completed.data # type: ignore[assignment]
assert completed
assert output is not None
msgs: list[ChatMessage] = output
# Expect: [user, A1 reply, summary]
assert len(msgs) == 3
assert msgs[0].role == Role.USER
@@ -17,7 +17,7 @@ from agent_framework._workflow._edge import (
SwitchCaseEdgeGroupCase,
SwitchCaseEdgeGroupDefault,
)
from agent_framework._workflow._executor import (
from agent_framework._workflow._workflow_executor import (
WorkflowExecutor,
)
@@ -3,6 +3,8 @@
import asyncio
from dataclasses import dataclass
from typing_extensions import Never
from agent_framework import (
Executor,
WorkflowBuilder,
@@ -33,13 +35,11 @@ class SimpleSubExecutor(Executor):
super().__init__(id="simple_sub")
@handler
async def process(self, request: SimpleRequest, ctx: WorkflowContext[None]) -> None:
async def process(self, request: SimpleRequest, ctx: WorkflowContext[Never, SimpleResponse]) -> None:
"""Process a simple request."""
from agent_framework import WorkflowCompletedEvent
# Just echo back with prefix and complete
response = SimpleResponse(result=f"processed: {request.text}")
await ctx.add_event(WorkflowCompletedEvent(data=response))
await ctx.yield_output(response)
class SimpleParent(Executor):
@@ -57,7 +57,7 @@ class SimpleParent(Executor):
await ctx.send_message(request, target_id="sub_workflow")
@handler
async def collect(self, response: SimpleResponse, ctx: WorkflowContext[None]) -> None:
async def collect(self, response: SimpleResponse, ctx: WorkflowContext) -> None:
"""Collect the result."""
self.result = response
@@ -72,7 +72,7 @@ async def test_simple_sub_workflow():
super().__init__(id="dummy")
@handler
async def process(self, message: object, ctx: WorkflowContext[None]) -> None:
async def process(self, message: object, ctx: WorkflowContext) -> None:
pass # Do nothing
dummy = DummyExecutor()
@@ -5,6 +5,7 @@ from dataclasses import dataclass
from typing import Any
from pydantic import Field
from typing_extensions import Never
from agent_framework import (
Executor,
@@ -12,7 +13,6 @@ from agent_framework import (
RequestInfoMessage,
RequestResponse,
WorkflowBuilder,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowExecutor,
handler,
@@ -54,7 +54,7 @@ class EmailValidator(Executor):
@handler
async def validate_request(
self, request: EmailValidationRequest, ctx: WorkflowContext[RequestInfoMessage | ValidationResult]
self, request: EmailValidationRequest, ctx: WorkflowContext[RequestInfoMessage, ValidationResult]
) -> None:
"""Validate an email address."""
# Extract domain and check if it's approved
@@ -62,7 +62,7 @@ class EmailValidator(Executor):
if not domain:
result = ValidationResult(email=request.email, is_valid=False, reason="Invalid email format")
await ctx.add_event(WorkflowCompletedEvent(data=result))
await ctx.yield_output(result)
return
# Request domain check from external source
@@ -71,7 +71,7 @@ class EmailValidator(Executor):
@handler
async def handle_domain_response(
self, response: RequestResponse[DomainCheckRequest, bool], ctx: WorkflowContext[ValidationResult]
self, response: RequestResponse[DomainCheckRequest, bool], ctx: WorkflowContext[Never, ValidationResult]
) -> None:
"""Handle domain check response with correlation."""
# Use the original email from the correlated response
@@ -80,7 +80,7 @@ class EmailValidator(Executor):
is_valid=response.data or False,
reason="Domain approved" if response.data else "Domain not approved",
)
await ctx.add_event(WorkflowCompletedEvent(data=result))
await ctx.yield_output(result)
class ParentOrchestrator(Executor):
@@ -114,7 +114,7 @@ class ParentOrchestrator(Executor):
return RequestResponse[DomainCheckRequest, bool].forward()
@handler
async def collect_result(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None:
async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None:
"""Collect validation results."""
self.results.append(result)
@@ -146,11 +146,11 @@ async def test_basic_sub_workflow() -> None:
await ctx.send_message(request, target_id="email_workflow")
@handler
async def collect(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None:
async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None:
self.result = result
parent = SimpleParent()
workflow_executor = WorkflowExecutor(validation_workflow, id="email_workflow")
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
main_request_info = RequestInfoExecutor(id="main_request_info")
main_workflow = (
@@ -199,7 +199,7 @@ async def test_sub_workflow_with_interception():
# Create parent workflow with interception
parent = ParentOrchestrator(approved_domains={"example.com", "internal.org"})
workflow_executor = WorkflowExecutor(validation_workflow, id="email_workflow")
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
parent_request_info = RequestInfoExecutor(id="request_info")
main_workflow = (
@@ -275,7 +275,7 @@ async def test_conditional_forwarding() -> None:
return RequestResponse[DomainCheckRequest, bool].forward()
@handler
async def collect(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None:
async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None:
self.result = result
# Setup workflows
@@ -291,7 +291,7 @@ async def test_conditional_forwarding() -> None:
)
parent = ConditionalParent()
workflow_executor = WorkflowExecutor(validation_workflow, id="email_workflow")
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
parent_request_info = RequestInfoExecutor(id="request_info")
main_workflow = (
@@ -358,7 +358,7 @@ async def test_workflow_scoped_interception() -> None:
return RequestResponse[DomainCheckRequest, bool].forward()
@handler
async def collect(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None:
async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None:
self.results[result.email] = result
# Create two identical sub-workflows
@@ -377,8 +377,8 @@ async def test_workflow_scoped_interception() -> None:
workflow_b = create_validation_workflow()
parent = MultiWorkflowParent()
executor_a = WorkflowExecutor(workflow_a, id="workflow_a")
executor_b = WorkflowExecutor(workflow_b, id="workflow_b")
executor_a = WorkflowExecutor(workflow_a, "workflow_a")
executor_b = WorkflowExecutor(workflow_b, "workflow_b")
parent_request_info = RequestInfoExecutor(id="request_info")
main_workflow = (
@@ -407,9 +407,102 @@ async def test_workflow_scoped_interception() -> None:
assert parent.results["user@random.com"].is_valid is True
async def test_concurrent_sub_workflow_execution() -> None:
"""Test that WorkflowExecutor can handle multiple concurrent invocations properly."""
class ConcurrentProcessor(Executor):
"""Processor that sends multiple concurrent requests to the same sub-workflow."""
results: list[ValidationResult] = Field(default_factory=list)
def __init__(self, **kwargs: Any):
super().__init__(id="concurrent_processor", **kwargs)
@handler
async def start(self, emails: list[str], ctx: WorkflowContext[EmailValidationRequest]) -> None:
"""Send multiple concurrent requests to the same sub-workflow."""
# Send all requests concurrently to the same workflow executor
for email in emails:
request = EmailValidationRequest(email=email)
await ctx.send_message(request, target_id="email_workflow")
@handler
async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None:
"""Collect results from concurrent executions."""
self.results.append(result)
# Create sub-workflow for email validation
email_validator = EmailValidator()
email_request_info = RequestInfoExecutor(id="email_request_info")
validation_workflow = (
WorkflowBuilder()
.set_start_executor(email_validator)
.add_edge(email_validator, email_request_info)
.add_edge(email_request_info, email_validator)
.build()
)
# Create parent workflow
processor = ConcurrentProcessor()
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
parent_request_info = RequestInfoExecutor(id="request_info")
main_workflow = (
WorkflowBuilder()
.set_start_executor(processor)
.add_edge(processor, workflow_executor)
.add_edge(workflow_executor, processor)
.add_edge(workflow_executor, parent_request_info) # For external requests
.add_edge(parent_request_info, workflow_executor) # For SubWorkflowResponse routing
.build()
)
# Test concurrent execution with multiple emails
emails = [
"user1@domain1.com",
"user2@domain2.com",
"user3@domain3.com",
"user4@domain4.com",
"user5@domain5.com",
]
result = await main_workflow.run(emails)
# Each email should generate one external request
request_events = result.get_request_info_events()
assert len(request_events) == len(emails)
# Verify each request corresponds to the correct domain
domains_requested = {event.data.domain for event in request_events} # type: ignore[union-attr]
expected_domains = {f"domain{i}.com" for i in range(1, 6)}
assert domains_requested == expected_domains
# Send responses for all requests (approve all domains)
responses = {event.request_id: True for event in request_events}
await main_workflow.send_responses(responses)
# All results should be collected
assert len(processor.results) == len(emails)
# Verify each email was processed correctly
result_emails = {result.email for result in processor.results}
expected_emails = set(emails)
assert result_emails == expected_emails
# All should be valid since we approved all domains
for result_obj in processor.results:
assert result_obj.is_valid is True
assert result_obj.reason == "Domain approved"
# Verify that concurrent executions were properly isolated
# (This is implicitly tested by the fact that we got correct results for all emails)
if __name__ == "__main__":
# Run tests
asyncio.run(test_basic_sub_workflow())
asyncio.run(test_sub_workflow_with_interception())
asyncio.run(test_conditional_forwarding())
asyncio.run(test_workflow_scoped_interception())
asyncio.run(test_concurrent_sub_workflow_execution())
@@ -19,7 +19,6 @@ from agent_framework import (
validate_workflow_graph,
)
from agent_framework._workflow._edge import SingleEdgeGroup
from agent_framework._workflow._validation import HandlerOutputAnnotationError
class StringExecutor(Executor):
@@ -51,8 +50,8 @@ class AnyExecutor(Executor):
class NoOutputTypesExecutor(Executor):
@handler
async def handle_message(self, message: str, ctx: WorkflowContext[Any]) -> None:
await ctx.send_message("processed")
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
await ctx.send_message("processed") # type: ignore[arg-type]
class MultiTypeExecutor(Executor):
@@ -575,58 +574,33 @@ def test_validation_enum_usage() -> None:
def test_handler_ctx_missing_annotation_raises() -> None:
class BadExecutor(Executor):
@handler
async def handle(self, message: str, ctx) -> None: # type: ignore[no-untyped-def]
pass
# Validation now happens at handler registration time, not workflow build time
with pytest.raises(ValueError) as exc:
start = StringExecutor(id="s")
bad = BadExecutor(id="b")
class BadExecutor(Executor):
@handler
async def handle(self, message: str, ctx) -> None: # type: ignore[no-untyped-def]
pass
with pytest.raises(HandlerOutputAnnotationError) as exc:
WorkflowBuilder().add_edge(start, bad).set_start_executor(start).build()
assert exc.value.validation_type == ValidationTypeEnum.HANDLER_OUTPUT_ANNOTATION
assert "missing type annotation" in str(exc.value)
def test_handler_ctx_unsubscripted_workflow_context_raises() -> None:
class BadExecutor(Executor):
@handler
async def handle(self, message: str, ctx: WorkflowContext) -> None: # type: ignore # missing T
pass
start = StringExecutor(id="s")
bad = BadExecutor(id="b")
with pytest.raises(HandlerOutputAnnotationError) as exc:
WorkflowBuilder().add_edge(start, bad).set_start_executor(start).build()
assert exc.value.validation_type == ValidationTypeEnum.HANDLER_OUTPUT_ANNOTATION
# Message should mention missing T or WorkflowContext[None]
assert "WorkflowContext[None]" in str(exc.value) or "missing" in str(exc.value).lower()
assert "must have a WorkflowContext" in str(exc.value)
def test_handler_ctx_invalid_t_out_entries_raises() -> None:
class BadExecutor(Executor):
@handler
async def handle(self, message: str, ctx: WorkflowContext[123]) -> None: # type: ignore[valid-type]
pass
# Validation now happens at handler registration time, not workflow build time
with pytest.raises(ValueError) as exc:
start = StringExecutor(id="s")
bad = BadExecutor(id="b")
class BadExecutor(Executor):
@handler
async def handle(self, message: str, ctx: WorkflowContext[123]) -> None: # type: ignore[valid-type]
pass
with pytest.raises(HandlerOutputAnnotationError) as exc:
WorkflowBuilder().add_edge(start, bad).set_start_executor(start).build()
assert exc.value.validation_type == ValidationTypeEnum.HANDLER_OUTPUT_ANNOTATION
assert "invalid entries" in str(exc.value)
assert "invalid type entry" in str(exc.value)
def test_handler_ctx_none_is_allowed() -> None:
class NoneExecutor(Executor):
@handler
async def handle(self, message: str, ctx: WorkflowContext[None]) -> None:
async def handle(self, message: str, ctx: WorkflowContext) -> None:
# does not emit
return None
@@ -11,7 +11,7 @@ class MockExecutor(Executor):
"""A mock executor for testing purposes."""
@handler
async def mock_handler(self, message: str, ctx: WorkflowContext[None]) -> None:
async def mock_handler(self, message: str, ctx: WorkflowContext) -> None:
"""A mock handler that does nothing."""
pass
@@ -20,7 +20,7 @@ class ListStrTargetExecutor(Executor):
"""A mock executor that accepts a list of strings (for fan-in targets)."""
@handler
async def handle(self, message: list[str], ctx: WorkflowContext[None]) -> None: # type: ignore[type-arg]
async def handle(self, message: list[str], ctx: WorkflowContext) -> None:
pass
@@ -15,9 +15,11 @@ from agent_framework import (
RequestInfoMessage,
RequestResponse,
WorkflowBuilder,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowEvent,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
@@ -36,20 +38,20 @@ class IncrementExecutor(Executor):
increment: int = 1
@handler
async def mock_handler(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage]) -> None:
async def mock_handler(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage, int]) -> None:
if message.data < self.limit:
await ctx.send_message(NumberMessage(data=message.data + self.increment))
else:
await ctx.add_event(WorkflowCompletedEvent(data=message.data))
await ctx.yield_output(message.data)
class AggregatorExecutor(Executor):
"""A mock executor that aggregates results from multiple executors."""
@handler
async def mock_handler(self, messages: list[NumberMessage], ctx: WorkflowContext[Any]) -> None:
# This mock simply returns the data incremented by 1
await ctx.add_event(WorkflowCompletedEvent(data=sum(msg.data for msg in messages)))
async def mock_handler(self, messages: list[NumberMessage], ctx: WorkflowContext[Any, int]) -> None:
# This mock simply returns the sum of the data
await ctx.yield_output(sum(msg.data for msg in messages))
@dataclass
@@ -70,18 +72,21 @@ class MockExecutorRequestApproval(Executor):
@handler
async def mock_handler_b(
self, message: RequestResponse[RequestInfoMessage, ApprovalMessage], ctx: WorkflowContext[NumberMessage]
self,
message: RequestResponse[RequestInfoMessage, ApprovalMessage],
ctx: WorkflowContext[NumberMessage, int],
) -> None:
"""A mock handler that processes the approval response."""
data = await ctx.get_shared_state(self.id)
assert isinstance(data, int)
assert isinstance(message.data, ApprovalMessage)
if message.data.approved:
await ctx.add_event(WorkflowCompletedEvent(data=data))
await ctx.yield_output(data)
else:
await ctx.send_message(NumberMessage(data=data))
async def test_workflow_run_streaming():
async def test_workflow_run_streaming() -> None:
"""Test the workflow run stream."""
executor_a = IncrementExecutor(id="executor_a")
executor_b = IncrementExecutor(id="executor_b")
@@ -97,7 +102,7 @@ async def test_workflow_run_streaming():
result: int | None = None
async for event in workflow.run_stream(NumberMessage(data=0)):
assert isinstance(event, WorkflowEvent)
if isinstance(event, WorkflowCompletedEvent):
if isinstance(event, WorkflowOutputEvent):
result = event.data
assert result is not None and result == 10
@@ -136,9 +141,9 @@ async def test_workflow_run():
)
events = await workflow.run(NumberMessage(data=0))
completed_event = events.get_completed_event()
assert isinstance(completed_event, WorkflowCompletedEvent)
assert completed_event.data == 10
assert events.get_final_state() == WorkflowRunState.IDLE
outputs = events.get_outputs()
assert outputs[0] == 10
async def test_workflow_run_not_completed():
@@ -182,13 +187,18 @@ async def test_workflow_send_responses_streaming():
assert request_info_event is not None
result: int | None = None
completed = False
async for event in workflow.send_responses_streaming({
request_info_event.request_id: ApprovalMessage(approved=True)
}):
if isinstance(event, WorkflowCompletedEvent):
if isinstance(event, WorkflowOutputEvent):
result = event.data
elif isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
completed = True
assert result is not None and result == 1 # The data should be incremented by 1 from the initial message
assert (
completed and result is not None and result == 1
) # The data should be incremented by 1 from the initial message
async def test_workflow_send_responses():
@@ -214,9 +224,9 @@ async def test_workflow_send_responses():
result = await workflow.send_responses({request_info_events[0].request_id: ApprovalMessage(approved=True)})
completed_event = result.get_completed_event()
assert isinstance(completed_event, WorkflowCompletedEvent)
assert completed_event.data == 1 # The data should be incremented by 1 from the initial message
assert result.get_final_state() == WorkflowRunState.IDLE
outputs = result.get_outputs()
assert outputs[0] == 1 # The data should be incremented by 1 from the initial message
async def test_fan_out():
@@ -232,11 +242,12 @@ async def test_fan_out():
events = await workflow.run(NumberMessage(data=0))
# Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent
# executor_b will also emit a WorkflowCompletedEvent
# executor_b will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore)
assert len(events) == 7
completed_event = events.get_completed_event()
assert completed_event is not None and completed_event.data == 1
assert events.get_final_state() == WorkflowRunState.IDLE
outputs = events.get_outputs()
assert outputs[0] == 1
async def test_fan_out_multiple_completed_events():
@@ -252,11 +263,12 @@ async def test_fan_out_multiple_completed_events():
events = await workflow.run(NumberMessage(data=0))
# Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent
# executor_a and executor_b will also emit a WorkflowCompletedEvent
# executor_b and executor_c will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore)
assert len(events) == 8
with pytest.raises(ValueError):
events.get_completed_event()
# Multiple outputs are expected from both executors
outputs = events.get_outputs()
assert len(outputs) == 2
async def test_fan_in():
@@ -277,18 +289,19 @@ async def test_fan_in():
events = await workflow.run(NumberMessage(data=0))
# Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent
# aggregator will also emit a WorkflowCompletedEvent
# aggregator will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore)
assert len(events) == 9
completed_event = events.get_completed_event()
assert completed_event is not None and completed_event.data == 4
assert events.get_final_state() == WorkflowRunState.IDLE
outputs = events.get_outputs()
assert outputs[0] == 4 # executor_a(0->1), both executor_b and executor_c(1->2), aggregator(2+2=4)
@pytest.fixture
def simple_executor() -> Executor:
class SimpleExecutor(Executor):
@handler
async def handle_message(self, message: str, context: WorkflowContext[None]) -> None:
async def handle_message(self, message: str, context: WorkflowContext) -> None:
pass
return SimpleExecutor(id="test_executor")
@@ -442,7 +455,7 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu
# Test non-streaming run_from_checkpoint method
result = await workflow.run_from_checkpoint(checkpoint_id)
assert isinstance(result, list) # Should return WorkflowRunResult which extends list
assert hasattr(result, "get_completed_event") # Should have WorkflowRunResult methods
assert hasattr(result, "get_outputs") # Should have WorkflowRunResult methods
async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executor: Executor):
@@ -498,7 +511,7 @@ class StateTrackingExecutor(Executor):
"""An executor that tracks state in shared state to test context reset behavior."""
@handler
async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any]) -> None:
async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any, list]) -> None:
"""Handle the message and track it in shared state."""
# Get existing messages from shared state
try:
@@ -513,8 +526,8 @@ class StateTrackingExecutor(Executor):
# Update shared state
await ctx.set_shared_state("processed_messages", existing_messages)
# Complete workflow with current shared state
await ctx.add_event(WorkflowCompletedEvent(data=existing_messages.copy())) # type: ignore
# Yield output
await ctx.yield_output(existing_messages.copy()) # type: ignore
async def test_workflow_multiple_runs_no_state_collision():
@@ -536,27 +549,27 @@ async def test_workflow_multiple_runs_no_state_collision():
# Run 1: Should only see messages from run 1
result1 = await workflow.run(StateTrackingMessage(data="message1", run_id="run1"))
completed1 = result1.get_completed_event()
assert completed1 is not None
assert completed1.data == ["run1:message1"]
assert result1.get_final_state() == WorkflowRunState.IDLE
outputs1 = result1.get_outputs()
assert outputs1[0] == ["run1:message1"]
# Run 2: Should only see messages from run 2, not run 1
result2 = await workflow.run(StateTrackingMessage(data="message2", run_id="run2"))
completed2 = result2.get_completed_event()
assert completed2 is not None
assert completed2.data == ["run2:message2"] # Should NOT contain run1 data
assert result2.get_final_state() == WorkflowRunState.IDLE
outputs2 = result2.get_outputs()
assert outputs2[0] == ["run2:message2"] # Should NOT contain run1 data
# Run 3: Should only see messages from run 3
result3 = await workflow.run(StateTrackingMessage(data="message3", run_id="run3"))
completed3 = result3.get_completed_event()
assert completed3 is not None
assert completed3.data == ["run3:message3"] # Should NOT contain run1 or run2 data
assert result3.get_final_state() == WorkflowRunState.IDLE
outputs3 = result3.get_outputs()
assert outputs3[0] == ["run3:message3"] # Should NOT contain run1 or run2 data
# Verify that each run only processed its own message
# This confirms that the checkpointable context properly resets between runs
assert completed1.data != completed2.data
assert completed2.data != completed3.data
assert completed1.data != completed3.data
assert outputs1[0] != outputs2[0]
assert outputs2[0] != outputs3[0]
assert outputs1[0] != outputs3[0]
async def test_comprehensive_edge_groups_workflow():
@@ -604,17 +617,17 @@ async def test_comprehensive_edge_groups_workflow():
# router(2->3) -> switch routes to proc_a -> proc_a(3->4) -> fanout_hub(4->5)
# -> [parallel_1(5->8), parallel_2(5->10)] -> aggregator(8+10=18)
events_small = await workflow.run(NumberMessage(data=2))
completed_small = events_small.get_completed_event()
assert completed_small is not None
assert completed_small.data == 18 # Exact expected result: 8+10 from parallel processors
assert events_small.get_final_state() == WorkflowRunState.IDLE
outputs_small = events_small.get_outputs()
assert outputs_small[0] == 18 # Exact expected result: 8+10 from parallel processors
# Test with large number (should go through processor_b)
# router(8->9) -> switch routes to proc_b -> proc_b(9->11) -> fanout_hub(11->12)
# -> [parallel_1(12->15), parallel_2(12->17)] -> aggregator(15+17=32)
events_large = await workflow.run(NumberMessage(data=8))
completed_large = events_large.get_completed_event()
assert completed_large is not None
assert completed_large.data == 32 # Exact expected result: 15+17 from parallel processors
assert events_large.get_final_state() == WorkflowRunState.IDLE
outputs_large = events_large.get_outputs()
assert outputs_large[0] == 32 # Exact expected result: 15+17 from parallel processors
# The key verification is that we successfully executed a workflow using all three edge group types
# and that both switch-case paths work (small vs large numbers)
@@ -624,9 +637,9 @@ async def test_comprehensive_edge_groups_workflow():
assert len(events_large) >= 6
# Verify different paths were taken by checking exact results
assert completed_small.data == 18, f"Small number path should result in 18, got {completed_small.data}"
assert completed_large.data == 32, f"Large number path should result in 32, got {completed_large.data}"
assert completed_small.data != completed_large.data, "Different paths should produce different results"
assert outputs_small[0] == 18, f"Small number path should result in 18, got {outputs_small[0]}"
assert outputs_large[0] == 32, f"Large number path should result in 32, got {outputs_large[0]}"
assert outputs_small[0] != outputs_large[0], "Different paths should produce different results"
# Both tests should complete successfully, proving all edge group types work
@@ -660,11 +673,9 @@ async def test_workflow_with_simple_cycle_and_exit_condition():
# Test the cycle
# Expected: exec_a(2->4) -> exec_b(4->5) -> exec_a(5->7, completes because 7 >= 6)
events = await workflow.run(NumberMessage(data=2))
completed_event = events.get_completed_event()
assert completed_event is not None
assert (
completed_event.data is not None and completed_event.data >= 6
) # Should complete when executor_a reaches its limit
assert events.get_final_state() == WorkflowRunState.IDLE
outputs = events.get_outputs()
assert outputs[0] is not None and outputs[0] >= 6 # Should complete when executor_a reaches its limit
# Verify cycling occurred (should have events from both executors)
# Check for ExecutorInvokedEvent and ExecutorCompletedEvent types that have executor_id
@@ -3,14 +3,19 @@
import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
from typing_extensions import Never
from agent_framework import (
WorkflowCompletedEvent,
Executor,
WorkflowBuilder,
WorkflowContext,
WorkflowEvent,
WorkflowRunState,
WorkflowStatusEvent,
executor,
handler,
)
if TYPE_CHECKING:
@@ -57,8 +62,203 @@ async def test_executor_cannot_emit_framework_lifecycle_event(caplog: "LogCaptur
async def test_executor_emits_normal_event() -> None:
async with make_context() as (ctx, runner_ctx):
await ctx.add_event(WorkflowCompletedEvent("done"))
# Create a normal event to test event emission
await ctx.add_event(_TestEvent())
events: list[WorkflowEvent] = await runner_ctx.drain_events()
assert len(events) == 1
assert isinstance(events[0], WorkflowCompletedEvent)
assert isinstance(events[0], _TestEvent)
class _TestEvent(WorkflowEvent):
pass
async def test_workflow_context_type_annotations_no_parameter() -> None:
# Test function-based executor
@executor(id="func1")
async def func1(text: str, ctx: WorkflowContext) -> None:
await ctx.add_event(_TestEvent())
wf = WorkflowBuilder().set_start_executor(func1).build()
events = await wf.run("hello")
test_events = [e for e in events if isinstance(e, _TestEvent)]
assert len(test_events) == 1
# Test class-based executor
class _exec1(Executor):
@handler
async def func1(self, text: str, ctx: WorkflowContext) -> None:
await ctx.add_event(_TestEvent())
executor1 = _exec1(id="exec1")
assert executor1.input_types == [str]
assert executor1.output_types == []
assert executor1.workflow_output_types == []
wf2 = WorkflowBuilder().set_start_executor(executor1).build()
events2 = await wf2.run("hello")
test_events2 = [e for e in events2 if isinstance(e, _TestEvent)]
assert len(test_events2) == 1
async def test_workflow_context_type_annotations_message_type_parameter() -> None:
# Test function-based executor
@executor(id="func1")
async def func1(text: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message("world")
@executor(id="func2")
async def func2(text: str, ctx: WorkflowContext) -> None:
await ctx.add_event(_TestEvent(data=text))
wf = WorkflowBuilder().add_edge(func1, func2).set_start_executor(func1).build()
events = await wf.run("hello")
test_events = [e for e in events if isinstance(e, _TestEvent)]
assert len(test_events) == 1
assert test_events[0].data == "world"
# Test class-based executor
class _exec1(Executor):
@handler
async def func1(self, text: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message("world")
class _exec2(Executor):
@handler
async def func2(self, text: str, ctx: WorkflowContext) -> None:
await ctx.add_event(_TestEvent(data=text))
executor1 = _exec1(id="exec1")
executor2 = _exec2(id="exec2")
assert executor1.input_types == [str]
assert executor1.output_types == [str]
assert executor1.workflow_output_types == []
assert executor2.input_types == [str]
assert executor2.output_types == []
assert executor2.workflow_output_types == []
wf2 = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
events2 = await wf2.run("hello")
test_events2 = [e for e in events2 if isinstance(e, _TestEvent)]
assert len(test_events2) == 1
assert test_events2[0].data == "world"
async def test_workflow_context_type_annotations_message_and_output_type_parameters() -> None:
# Test function-based executor
@executor(id="func1")
async def func1(text: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message("world")
@executor(id="func2")
async def func2(text: str, ctx: WorkflowContext[Never, str]) -> None:
await ctx.add_event(_TestEvent(data=text))
await ctx.yield_output(text)
wf = WorkflowBuilder().add_edge(func1, func2).set_start_executor(func1).build()
events = await wf.run("hello")
outputs = events.get_outputs()
assert len(outputs) == 1
assert outputs[0] == "world"
# Test class-based executor
class _exec1(Executor):
@handler
async def func1(self, text: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message("world")
class _exec2(Executor):
@handler
async def func2(self, text: str, ctx: WorkflowContext[Never, str]) -> None:
await ctx.add_event(_TestEvent(data=text))
await ctx.yield_output(text)
executor1 = _exec1(id="exec1")
executor2 = _exec2(id="exec2")
assert executor1.input_types == [str]
assert executor1.output_types == [str]
assert executor1.workflow_output_types == []
assert executor2.input_types == [str]
assert executor2.output_types == []
assert executor2.workflow_output_types == [str]
wf2 = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
events2 = await wf2.run("hello")
outputs2 = events2.get_outputs()
assert len(outputs2) == 1
assert outputs2[0] == "world"
async def test_workflow_context_type_annotations_any() -> None:
class _exec1(Executor):
@handler
async def func1(self, text: str, ctx: WorkflowContext[Any]) -> None:
await ctx.add_event(_TestEvent())
await ctx.send_message(123)
executor1 = _exec1(id="exec1")
assert executor1.input_types == [str]
assert executor1.output_types == [Any]
class _exec2(Executor):
@handler
async def func2(self, number: int, ctx: WorkflowContext[Any, Any]) -> None:
await ctx.add_event(_TestEvent())
await ctx.send_message(456)
await ctx.yield_output(3.14)
executor2 = _exec2(id="exec2")
assert executor2.input_types == [int]
assert executor2.output_types == [Any]
assert executor2.workflow_output_types == [Any]
async def test_workflow_context_missing_annotation_error() -> None:
"""Test that missing WorkflowContext annotation raises appropriate error."""
import pytest
# Test function-based executor with missing ctx annotation
with pytest.raises(ValueError, match="must have a WorkflowContext"):
@executor(id="bad_func")
async def bad_func(text: str, ctx) -> None: # type: ignore[no-untyped-def]
pass
# Test class-based executor with missing ctx annotation
with pytest.raises(ValueError, match="must have a WorkflowContext"):
class _BadExecutor(Executor):
@handler
async def bad_handler(self, text: str, ctx) -> None: # type: ignore[no-untyped-def]
pass
async def test_workflow_context_invalid_type_parameter_error() -> None:
"""Test that invalid type parameters like int values raise appropriate errors."""
import pytest
# Test function-based executor with invalid type parameter (int value instead of type)
with pytest.raises(ValueError, match="invalid type entry"):
@executor(id="bad_func")
async def bad_func(text: str, ctx: WorkflowContext[123]) -> None: # type: ignore[valid-type]
pass
# Test class-based executor with invalid type parameter
with pytest.raises(ValueError, match="invalid type entry"):
class _BadExecutor(Executor):
@handler
async def bad_handler(self, text: str, ctx: WorkflowContext[456]) -> None: # type: ignore[valid-type]
pass
# Test two-parameter WorkflowContext with invalid workflow output type
with pytest.raises(ValueError, match="invalid type entry"):
@executor(id="bad_func2")
async def bad_func2(text: str, ctx: WorkflowContext[str, 789]) -> None: # type: ignore[valid-type]
pass
@@ -48,7 +48,7 @@ class SecondExecutor(Executor):
self._processed_messages: list[str] = []
@handler
async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None:
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
"""Handle string messages."""
self._processed_messages.append(message)
@@ -87,7 +87,7 @@ class FanInAggregator(Executor):
self._processed_messages: list[Any] = []
@handler
async def handle_aggregated_data(self, messages: list[str], ctx: WorkflowContext[None]) -> None:
async def handle_aggregated_data(self, messages: list[str], ctx: WorkflowContext) -> None:
# Process aggregated messages from fan-in
aggregated = f"aggregated: {', '.join(messages)}"
self._processed_messages.append(aggregated)
@@ -196,7 +196,14 @@ async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> No
assert message.source_span_id is not None
# Test executor trace context handling
await executor.execute("test message", workflow_ctx)
await executor.execute(
"test message",
["source"], # source_executor_ids
shared_state, # shared_state
ctx, # runner_context
trace_contexts=[{"traceparent": "00-12345678901234567890123456789012-1234567890123456-01"}],
source_span_ids=["1234567890123456"],
)
# Check that spans were created with proper attributes
spans = span_exporter.get_finished_spans()
@@ -372,7 +379,7 @@ async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExp
super().__init__(id="failing_executor")
@handler
async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None:
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
raise ValueError("Test error")
failing_executor = FailingExecutor()
@@ -4,6 +4,7 @@ from dataclasses import dataclass
from typing import Any
import pytest
from typing_extensions import Never
from agent_framework import (
Executor,
@@ -15,7 +16,6 @@ from agent_framework import (
SharedState,
Workflow,
WorkflowBuilder,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowEventSource,
WorkflowFailedEvent,
@@ -32,7 +32,7 @@ class FailingExecutor(Executor):
"""Executor that raises at runtime to test failure signaling."""
@handler
async def fail(self, msg: int, ctx: WorkflowContext[None]) -> None: # pragma: no cover - invoked via workflow
async def fail(self, msg: int, ctx: WorkflowContext) -> None: # pragma: no cover - invoked via workflow
raise RuntimeError("boom")
@@ -57,14 +57,14 @@ async def test_executor_failed_and_workflow_failed_events_streaming():
async def test_executor_failed_event_emitted_on_direct_execute():
failing = FailingExecutor(id="f")
ctx = InProcRunnerContext()
wf_ctx: WFContext[None] = WFContext(
executor_id=failing.id,
source_executor_ids=["START"],
shared_state=SharedState(),
runner_context=ctx,
)
shared_state = SharedState()
with pytest.raises(RuntimeError, match="boom"):
await failing.execute(0, wf_ctx)
await failing.execute(
0,
["START"],
shared_state,
ctx,
)
drained = await ctx.drain_events()
failed = [e for e in drained if isinstance(e, ExecutorFailedEvent)]
assert failed
@@ -98,17 +98,17 @@ class Completer(Executor):
"""Executor that completes immediately with provided data for testing."""
@handler
async def run(self, msg: str, ctx: WorkflowContext[str]) -> None: # pragma: no cover
await ctx.add_event(WorkflowCompletedEvent(msg))
async def run(self, msg: str, ctx: WorkflowContext[Never, str]) -> None: # pragma: no cover
await ctx.yield_output(msg)
async def test_completed_status_streaming():
c = Completer(id="c")
wf = WorkflowBuilder().set_start_executor(c).build()
events = [ev async for ev in wf.run_stream("ok")] # no raise
# Last status should be COMPLETED
# Last status should be IDLE
status = [e for e in events if isinstance(e, WorkflowStatusEvent)]
assert status and status[-1].state == WorkflowRunState.COMPLETED
assert status and status[-1].state == WorkflowRunState.IDLE
assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status)
@@ -120,8 +120,12 @@ async def test_started_and_completed_event_origins():
started = next(e for e in events if isinstance(e, WorkflowStartedEvent))
assert started.origin is WorkflowEventSource.FRAMEWORK
completed = next(e for e in events if isinstance(e, WorkflowCompletedEvent))
assert completed.origin is WorkflowEventSource.EXECUTOR
# Check for IDLE status indicating completion
idle_status = next(
(e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE), None
)
assert idle_status is not None
assert idle_status.origin is WorkflowEventSource.FRAMEWORK
async def test_non_streaming_final_state_helpers():
@@ -129,7 +133,7 @@ async def test_non_streaming_final_state_helpers():
c = Completer(id="c")
wf1 = WorkflowBuilder().set_start_executor(c).build()
result1: WorkflowRunResult = await wf1.run("done")
assert result1.get_final_state() == WorkflowRunState.COMPLETED
assert result1.get_final_state() == WorkflowRunState.IDLE
# Idle-with-pending-request case
req = Requester(id="req")
@@ -145,7 +149,7 @@ async def test_run_includes_status_events_completed():
result: WorkflowRunResult = await wf.run("ok")
timeline = result.status_timeline()
assert timeline, "Expected status timeline in non-streaming run() results"
assert timeline[-1].state == WorkflowRunState.COMPLETED
assert timeline[-1].state == WorkflowRunState.IDLE
async def test_run_includes_status_events_idle_with_requests():