[BREAKING] Python: Add context mode to AgentExecutor (#4668)

* Add context mode to AgentExecutor

* Fix unit tests

* Address comments

* Address comments

* REvise context mode and add tests

* Add chain config to sequential builder

* Add sample

* Fix pipeline

* Address comments

* Address comments

python/packages/core/tests/workflow/test_agent_executor.py

@@ -16,12 +16,12 @@ from agent_framework import (
     Content,
     Message,
     ResponseStream,
+    WorkflowBuilder,
     WorkflowEvent,
     WorkflowRunState,
 )
 from agent_framework._workflows._agent_executor import AgentExecutorResponse
 from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
-from agent_framework.orchestrations import SequentialBuilder
 
 if TYPE_CHECKING:
     from _pytest.logging import LogCaptureFixture
@@ -139,7 +139,7 @@ async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks()
     """AgentExecutor should call get_final_response() so stream result hooks execute."""
     agent = _StreamingHookAgent(id="hook_agent", name="HookAgent")
     executor = AgentExecutor(agent, id="hook_exec")
-    workflow = SequentialBuilder(participants=[executor]).build()
+    workflow = WorkflowBuilder(start_executor=executor).build()
 
     output_events: list[Any] = []
     async for event in workflow.run("run hook test", stream=True):
@@ -154,8 +154,9 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     """Test that workflow checkpoint stores AgentExecutor's cache and session states and restores them correctly."""
     storage = InMemoryCheckpointStorage()
 
-    # Create initial agent with a custom session
-    initial_agent = _CountingAgent(id="test_agent", name="TestAgent")
+    # Create two agents to form a two-step workflow
+    initial_agent_a = _CountingAgent(id="agent_a", name="AgentA")
+    initial_agent_b = _CountingAgent(id="agent_b", name="AgentB")
     initial_session = AgentSession()
 
     # Add some initial messages to the session state to verify session state persistence
@@ -165,11 +166,12 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     ]
     initial_session.state["history"] = {"messages": initial_messages}
 
-    # Create AgentExecutor with the session
-    executor = AgentExecutor(initial_agent, session=initial_session)
+    # Create AgentExecutors — first executor gets the custom session
+    exec_a = AgentExecutor(initial_agent_a, id="exec_a", session=initial_session)
+    exec_b = AgentExecutor(initial_agent_b, id="exec_b")
 
-    # Build workflow with checkpointing enabled
-    wf = SequentialBuilder(participants=[executor], checkpoint_storage=storage).build()
+    # Build two-executor workflow with checkpointing enabled
+    wf = WorkflowBuilder(start_executor=exec_a, checkpoint_storage=storage).add_edge(exec_a, exec_b).build()
 
     # Run the workflow with a user message
     first_run_output: AgentExecutorResponse | None = None
@@ -180,27 +182,25 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
             break
 
     assert first_run_output is not None
-    assert initial_agent.call_count == 1
+    assert initial_agent_a.call_count == 1
 
     # Verify checkpoint was created
     checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
-    assert len(checkpoints) >= 2, (
-        "Expected at least 2 checkpoints. The first one is after the start executor, "
-        "and the second one is after the agent execution."
+    assert len(checkpoints) >= 2, "Expected at least 2 checkpoints: one after exec_a and one after exec_b."
+
+    # Get the first checkpoint that contains exec_a's state (taken after exec_a completes,
+    # before exec_b runs)
+    checkpoints.sort(key=lambda cp: cp.timestamp)
+    restore_checkpoint = next(
+        cp for cp in checkpoints if "_executor_state" in cp.state and "exec_a" in cp.state["_executor_state"]
     )
 
-    # Get the second checkpoint which should contain the state after processing
-    # the first message by the start executor in the sequential workflow
-    checkpoints.sort(key=lambda cp: cp.timestamp)
-    restore_checkpoint = checkpoints[1]
-
     # Verify checkpoint contains executor state with both cache and session
-    assert "_executor_state" in restore_checkpoint.state
     executor_states = restore_checkpoint.state["_executor_state"]
     assert isinstance(executor_states, dict)
-    assert executor.id in executor_states
+    assert exec_a.id in executor_states
 
-    executor_state = executor_states[executor.id]  # type: ignore[index]
+    executor_state = executor_states[exec_a.id]  # type: ignore[index]
     assert "cache" in executor_state, "Checkpoint should store executor cache state"
     assert "agent_session" in executor_state, "Checkpoint should store executor session state"
 
@@ -213,19 +213,26 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     assert "pending_agent_requests" in executor_state
     assert "pending_responses_to_agent" in executor_state
 
-    # Create a new agent and executor for restoration
+    # Create new agents and executors for restoration
     # This simulates starting from a fresh state and restoring from checkpoint
-    restored_agent = _CountingAgent(id="test_agent", name="TestAgent")
+    restored_agent_a = _CountingAgent(id="agent_a", name="AgentA")
+    restored_agent_b = _CountingAgent(id="agent_b", name="AgentB")
     restored_session = AgentSession()
-    restored_executor = AgentExecutor(restored_agent, session=restored_session)
+    restored_exec_a = AgentExecutor(restored_agent_a, id="exec_a", session=restored_session)
+    restored_exec_b = AgentExecutor(restored_agent_b, id="exec_b")
 
-    # Verify the restored agent starts with a fresh state
-    assert restored_agent.call_count == 0
+    # Verify the restored agents start with a fresh state
+    assert restored_agent_a.call_count == 0
+    assert restored_agent_b.call_count == 0
 
-    # Build new workflow with the restored executor
-    wf_resume = SequentialBuilder(participants=[restored_executor], checkpoint_storage=storage).build()
+    # Build new workflow with the restored executors
+    wf_resume = (
+        WorkflowBuilder(start_executor=restored_exec_a, checkpoint_storage=storage)
+        .add_edge(restored_exec_a, restored_exec_b)
+        .build()
+    )
 
-    # Resume from checkpoint
+    # Resume from checkpoint — exec_a already ran, so exec_b should run and produce output
     resumed_output: AgentExecutorResponse | None = None
     async for ev in wf_resume.run(checkpoint_id=restore_checkpoint.checkpoint_id, stream=True):
         if ev.type == "output":
@@ -239,7 +246,7 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     assert resumed_output is not None
 
     # Verify the restored executor's session state was restored
-    restored_session_obj = restored_executor._session  # type: ignore[reportPrivateUsage]
+    restored_session_obj = restored_exec_a._session  # type: ignore[reportPrivateUsage]
     assert restored_session_obj is not None
     assert restored_session_obj.session_id == initial_session.session_id
 
@@ -306,7 +313,7 @@ async def test_agent_executor_run_with_session_kwarg_does_not_raise() -> None:
     """Passing session= via workflow.run() should not cause a duplicate-keyword TypeError (#4295)."""
     agent = _CountingAgent(id="session_kwarg_agent", name="SessionKwargAgent")
     executor = AgentExecutor(agent, id="session_kwarg_exec")
-    workflow = SequentialBuilder(participants=[executor]).build()
+    workflow = WorkflowBuilder(start_executor=executor).build()
 
     # This previously raised: TypeError: run() got multiple values for keyword argument 'session'
     result = await workflow.run("hello", session="user-supplied-value")
@@ -318,7 +325,7 @@ async def test_agent_executor_run_streaming_with_stream_kwarg_does_not_raise() -
     """Passing stream= via workflow.run() kwargs should not cause a duplicate-keyword TypeError."""
     agent = _CountingAgent(id="stream_kwarg_agent", name="StreamKwargAgent")
     executor = AgentExecutor(agent, id="stream_kwarg_exec")
-    workflow = SequentialBuilder(participants=[executor]).build()
+    workflow = WorkflowBuilder(start_executor=executor).build()
 
     # stream=True at workflow level triggers streaming mode (returns async iterable)
     events: list[WorkflowEvent] = []
@@ -378,7 +385,7 @@ async def test_agent_executor_run_with_messages_kwarg_does_not_raise() -> None:
     """Passing messages= via workflow.run() kwargs should not cause a duplicate-keyword TypeError."""
     agent = _CountingAgent(id="messages_kwarg_agent", name="MessagesKwargAgent")
     executor = AgentExecutor(agent, id="messages_kwarg_exec")
-    workflow = SequentialBuilder(participants=[executor]).build()
+    workflow = WorkflowBuilder(start_executor=executor).build()
 
     result = await workflow.run("hello", messages=["stale"])
     assert result is not None
@@ -426,7 +433,7 @@ async def test_agent_executor_workflow_with_non_copyable_raw_representation() ->
     exec_a = AgentExecutor(agent_a, id="exec_a")
     exec_b = AgentExecutor(agent_b, id="exec_b")
 
-    workflow = SequentialBuilder(participants=[exec_a, exec_b]).build()
+    workflow = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
     events = await workflow.run("hello")
 
     completed = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
@@ -440,3 +447,194 @@ async def test_agent_executor_workflow_with_non_copyable_raw_representation() ->
     assert len(agent_responses) > 0
     assert agent_responses[0].text == "reply from AgentA"
     assert agent_responses[0].raw_representation is raw
+
+
+# ---------------------------------------------------------------------------
+# Context mode tests
+# ---------------------------------------------------------------------------
+
+
+class _MessageCapturingAgent(BaseAgent):
+    """Agent that records the messages it received and returns a configurable reply."""
+
+    def __init__(self, *, reply_text: str = "reply", **kwargs: Any):
+        super().__init__(**kwargs)
+        self.reply_text = reply_text
+        self.last_messages: list[Message] = []
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        captured: list[Message] = []
+        if messages:
+            for m in messages:  # type: ignore[union-attr]
+                if isinstance(m, Message):
+                    captured.append(m)
+                elif isinstance(m, str):
+                    captured.append(Message("user", [m]))
+        self.last_messages = captured
+
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(contents=[Content.from_text(text=self.reply_text)])
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", [self.reply_text])])
+
+        return _run()
+
+
+def test_context_mode_custom_requires_context_filter() -> None:
+    """context_mode='custom' without context_filter must raise ValueError."""
+    agent = _CountingAgent(id="a", name="A")
+    with pytest.raises(ValueError, match="context_filter must be provided"):
+        AgentExecutor(agent, context_mode="custom")
+
+
+def test_context_mode_custom_with_filter_succeeds() -> None:
+    """context_mode='custom' with a context_filter should not raise."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, context_mode="custom", context_filter=lambda msgs: msgs[-1:])
+    assert executor._context_mode == "custom"  # pyright: ignore[reportPrivateUsage]
+    assert executor._context_filter is not None  # pyright: ignore[reportPrivateUsage]
+
+
+def test_context_mode_defaults_to_full() -> None:
+    """Default context_mode should be 'full'."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent)
+    assert executor._context_mode == "full"  # pyright: ignore[reportPrivateUsage]
+
+
+def test_context_mode_invalid_value_raises() -> None:
+    """Invalid context_mode value should raise ValueError."""
+    agent = _CountingAgent(id="a", name="A")
+    with pytest.raises(ValueError, match="context_mode must be one of"):
+        AgentExecutor(agent, context_mode="invalid_mode")  # type: ignore
+
+
+async def test_from_response_context_mode_full_passes_full_conversation() -> None:
+    """context_mode='full' (default) should pass full_conversation to the second agent."""
+    first = _MessageCapturingAgent(id="first", name="First", reply_text="first reply")
+    second = _MessageCapturingAgent(id="second", name="Second", reply_text="second reply")
+
+    exec_a = AgentExecutor(first, id="exec_a")
+    exec_b = AgentExecutor(second, id="exec_b", context_mode="full")
+
+    wf = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
+
+    async for ev in wf.run("hello", stream=True):
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    # Second agent should see full conversation: [user("hello"), assistant("first reply")]
+    seen = second.last_messages
+    assert len(seen) == 2
+    assert seen[0].role == "user" and "hello" in (seen[0].text or "")
+    assert seen[1].role == "assistant" and "first reply" in (seen[1].text or "")
+
+
+async def test_from_response_context_mode_last_agent_passes_only_agent_messages() -> None:
+    """context_mode='last_agent' should pass only the previous agent's response messages."""
+    first = _MessageCapturingAgent(id="first", name="First", reply_text="first reply")
+    second = _MessageCapturingAgent(id="second", name="Second", reply_text="second reply")
+
+    exec_a = AgentExecutor(first, id="exec_a")
+    exec_b = AgentExecutor(second, id="exec_b", context_mode="last_agent")
+
+    wf = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
+
+    async for ev in wf.run("hello", stream=True):
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    # Second agent should see only the assistant message from first: [assistant("first reply")]
+    seen = second.last_messages
+    assert len(seen) == 1
+    assert seen[0].role == "assistant" and "first reply" in (seen[0].text or "")
+
+
+async def test_from_response_context_mode_custom_uses_filter() -> None:
+    """context_mode='custom' should invoke context_filter on full_conversation."""
+    first = _MessageCapturingAgent(id="first", name="First", reply_text="first reply")
+    second = _MessageCapturingAgent(id="second", name="Second", reply_text="second reply")
+
+    # Custom filter: keep only user messages
+    def only_user_messages(msgs: list[Message]) -> list[Message]:
+        return [m for m in msgs if m.role == "user"]
+
+    exec_a = AgentExecutor(first, id="exec_a")
+    exec_b = AgentExecutor(second, id="exec_b", context_mode="custom", context_filter=only_user_messages)
+
+    wf = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
+
+    async for ev in wf.run("hello", stream=True):
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    # Second agent should see only user messages: [user("hello")]
+    seen = second.last_messages
+    assert len(seen) == 1
+    assert seen[0].role == "user" and "hello" in (seen[0].text or "")
+
+
+async def test_checkpoint_save_does_not_include_context_mode() -> None:
+    """on_checkpoint_save should not include context_mode in the saved state."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, context_mode="last_agent")
+
+    state = await executor.on_checkpoint_save()
+
+    assert "context_mode" not in state
+    assert "cache" in state
+    assert "agent_session" in state
+
+
+async def test_checkpoint_restore_works_without_context_mode_in_state() -> None:
+    """on_checkpoint_restore should succeed when state does not contain context_mode."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, context_mode="last_agent")
+
+    # Simulate a checkpoint state without context_mode (as saved by the new code)
+    state: dict[str, Any] = {
+        "cache": [Message(role="user", text="cached msg")],
+        "full_conversation": [],
+        "agent_session": AgentSession().to_dict(),
+        "pending_agent_requests": {},
+        "pending_responses_to_agent": [],
+    }
+
+    await executor.on_checkpoint_restore(state)
+
+    cache = executor._cache  # pyright: ignore[reportPrivateUsage]
+    assert len(cache) == 1
+    assert cache[0].text == "cached msg"
+    # context_mode should remain as configured in the constructor, not changed by restore
+    assert executor._context_mode == "last_agent"  # pyright: ignore[reportPrivateUsage]

python/packages/core/tests/workflow/test_runner.py

@@ -341,7 +341,7 @@ async def test_runner_emits_runner_completion_for_agent_response_without_targets
 
     await ctx.send_message(
         WorkflowMessage(
-            data=AgentExecutorResponse("agent", AgentResponse()),
+            data=AgentExecutorResponse("agent", AgentResponse(), []),
             source_id="agent",
         )
     )