diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index dd4c48b6220..dfaa5555941 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -146,7 +146,9 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: # Propagate errors from iteration, but first surface any pending events try: await iteration_task - except Exception: + except (Exception, asyncio.CancelledError): + # Discard pending state writes from the failed or cancelled superstep + self._state.discard() # Make sure failure-related events (like ExecutorFailedEvent) are surfaced if await self._ctx.has_events(): for event in await self._ctx.drain_events(): @@ -224,11 +226,12 @@ async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: await gather_cancelling_siblings_on_error(*tasks) message_batches = await self._ctx.drain_messages() - tasks = [ - _deliver_messages(source_executor_id, source_messages) - for source_executor_id, source_messages in message_batches.items() - ] - await gather_cancelling_siblings_on_error(*tasks) + await gather_cancelling_siblings_on_error( + *( + _deliver_messages(source_executor_id, source_messages) + for source_executor_id, source_messages in message_batches.items() + ) + ) async def _prepare_checkpoint_state(self) -> None: """Persist executor snapshots into committed shared state. diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 3c924b4839f..92a15dafbd6 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import contextlib import gc import logging import tempfile @@ -659,6 +660,199 @@ def _build(): assert result2.get_outputs()[0] == ["run2:message2"] +@dataclass +class FlakyMessage: + """A message that can fail on demand for testing state discard behavior.""" + + fail: bool + + +class FlakyStateExecutor(Executor): + """An executor that fails on demand to test state discard on failure.""" + + @handler + async def handle_message( + self, + message: FlakyMessage, + ctx: WorkflowContext[FlakyMessage, str], + ) -> None: + if message.fail: + ctx.set_state("secret", "leaked-from-failed-run") + # Small delay to ensure cancellation can happen after write is staged + await asyncio.sleep(0.01) + raise RuntimeError("simulated transient failure") + + await ctx.yield_output("ok") + + +async def test_workflow_discards_pending_state_after_failed_superstep(): + """Test that pending state from a failed superstep is discarded and not committed. + + This is a regression test for GitHub issue #7859: pending state writes from + a failed superstep must not leak into a later successful run on the same + Workflow instance. + """ + workflow = WorkflowBuilder(start_executor=FlakyStateExecutor(id="flaky")).build() + + # First run: fails after staging a state write + with pytest.raises(RuntimeError, match="simulated transient failure"): + await workflow.run(FlakyMessage(fail=True)) + + # Verify the failed run did not leave the staged write pending + assert workflow._runner.state._pending == {} + + # Second run: succeeds without touching "secret" + result = await workflow.run(FlakyMessage(fail=False)) + assert result.get_final_state() == WorkflowRunState.IDLE + assert result.get_outputs() == ["ok"] + + # Verify the leaked state from the failed run is NOT in committed state + committed_state = workflow._runner.state.export_state() + assert "secret" not in committed_state + + +@dataclass +class FanOutTestMessage: + """Message for fan-out state leak test.""" + + should_fail: bool + + +class FanOutSourceExecutor(Executor): + """Source executor that sends fan-out test messages.""" + + @handler + async def handle(self, message: FanOutTestMessage, ctx: WorkflowContext[FanOutTestMessage]) -> None: + # Forward the message to targets + await ctx.send_message(message) + + +class FailingTargetExecutor(Executor): + """Target that fails immediately on execute.""" + + def __init__(self, id: str, b_started: asyncio.Event | None = None) -> None: + super().__init__(id=id) + self._b_started = b_started + + @handler + async def handle(self, message: FanOutTestMessage, ctx: WorkflowContext) -> None: + if message.should_fail: + # Wait for slow target to start before failing to ensure deterministic ordering + if self._b_started: + await self._b_started.wait() + raise RuntimeError("target A failed") + + +class SlowStateWritingTargetExecutor(Executor): + """Target that writes state after being unblocked (vulnerable pattern if not cancelled).""" + + def __init__(self, id: str, b_started: asyncio.Event, release_b: asyncio.Event, b_task_ref: list) -> None: + super().__init__(id=id) + self._b_started = b_started + self._release_b = release_b + self._b_task_ref = b_task_ref + + @handler + async def handle(self, message: FanOutTestMessage, ctx: WorkflowContext) -> None: + self._b_task_ref.append(asyncio.current_task()) + self._b_started.set() + if message.should_fail: + await self._release_b.wait() + ctx.set_state("leak_key", "leaked_value") + + +async def test_workflow_discards_pending_state_after_fanout_failure(): + """Test that pending state from a fan-out sibling target is discarded when another target fails. + + Regression test for GitHub issue #7859: when a fan-out superstep has one target fail + while a sibling target is still running, the sibling's state writes must not leak + into committed state. + """ + b_started = asyncio.Event() + release_b = asyncio.Event() + b_task_ref: list[asyncio.Task] = [] + + source = FanOutSourceExecutor(id="source") + failing_target = FailingTargetExecutor(id="failing_target", b_started=b_started) + slow_target = SlowStateWritingTargetExecutor( + id="slow_target", b_started=b_started, release_b=release_b, b_task_ref=b_task_ref + ) + + workflow = WorkflowBuilder(start_executor=source).add_fan_out_edges(source, [failing_target, slow_target]).build() + + # Verify topology: single FanOutEdgeGroup with two targets under one edge_runner + from agent_framework._workflows._edge import FanOutEdgeGroup + + fan_out_groups = [eg for eg in workflow.edge_groups if isinstance(eg, FanOutEdgeGroup)] + assert len(fan_out_groups) == 1, f"Expected 1 FanOutEdgeGroup, got {len(fan_out_groups)}" + assert fan_out_groups[0].target_ids == [failing_target.id, slow_target.id], ( + f"Expected targets [{failing_target.id}, {slow_target.id}], got {fan_out_groups[0].target_ids}" + ) + + async def _run_workflow() -> WorkflowRunResult: + return await workflow.run(FanOutTestMessage(should_fail=True)) + + run_task: asyncio.Task[WorkflowRunResult] = asyncio.create_task(_run_workflow()) + with pytest.raises(RuntimeError, match="target A failed"): + await asyncio.wait_for(run_task, timeout=5.0) + + release_b.set() + with contextlib.suppress(asyncio.CancelledError, Exception): + await asyncio.wait_for(b_task_ref[0], timeout=0.5) + + b_task_ref.clear() + + result = await workflow.run(FanOutTestMessage(should_fail=False)) + assert result.get_final_state() == WorkflowRunState.IDLE + + committed_state = workflow._runner.state.export_state() + assert "leak_key" not in committed_state + + +async def test_workflow_discards_pending_state_on_cancellation(): + """Test that pending state from a cancelled superstep is discarded and not committed. + + Regression test for PR #8819 review comment: asyncio.CancelledError inherits from + BaseException (not Exception since Python 3.8), so `except Exception:` never catches it. + When a superstep's task is cancelled, the discard() call must still run to prevent + pending writes from leaking into the next run. + + This test creates a workflow that will stage a state write, then yields control + before completing (via a slow executor). We cancel the task while it's mid-superstep + to ensure the write is staged and the runner's except block is hit. + """ + # Use the existing FlakyStateExecutor which stages a write then raises + workflow = WorkflowBuilder(start_executor=FlakyStateExecutor(id="flaky")).build() + + # Create a task that will stage a state write then fail + async def _run_failing(): + return await workflow.run(FlakyMessage(fail=True)) + + run_task = asyncio.create_task(_run_failing()) + + # Yield to let the task start and stage its write + await asyncio.sleep(0) + + # Cancel the task while it's mid-superstep (after write is staged but before error handling completes) + run_task.cancel() + + # Await the task, confirming CancelledError is raised (discard() must not swallow it) + with pytest.raises(asyncio.CancelledError): + await run_task + + # Verify the cancelled run did not leave the staged write pending + assert workflow._runner.state._pending == {} + + # Second run: succeeds without touching "secret" + result = await workflow.run(FlakyMessage(fail=False)) + assert result.get_final_state() == WorkflowRunState.IDLE + assert result.get_outputs() == ["ok"] + + # Verify the leaked state from the cancelled run is NOT in committed state + committed_state = workflow._runner.state.export_state() + assert "secret" not in committed_state + + async def test_workflow_checkpoint_runtime_only_configuration( simple_executor: Executor, ): diff --git a/python/packages/declarative/tests/test_http_request_executor.py b/python/packages/declarative/tests/test_http_request_executor.py index 1e71bd48b02..2d3002c7c79 100644 --- a/python/packages/declarative/tests/test_http_request_executor.py +++ b/python/packages/declarative/tests/test_http_request_executor.py @@ -584,13 +584,31 @@ async def test_response_headers_empty_assigned_none(self) -> None: @pytest.mark.asyncio async def test_non_2xx_still_publishes_headers(self) -> None: + """Non-2xx responses publish headers to pending state during the superstep. + + After PR #8819 fix: headers are written to pending state but are NOT + committed on error. The runner's discard() clears them along with other + pending writes from the failed superstep. This test verifies headers are + written to pending state during execution (visible via state.get() which + checks pending first), but confirms they are NOT durably persisted after + the error is raised and discard() runs. + + Note: This test checks committed state via export_state() after the error. + Before the fix, ctx.state.commit() in the error path would persist headers + to committed state. After the fix, discard() clears them from pending, + so they never reach committed state. + """ handler = StubHandler(_err(status=500, body="boom", headers={"X-Trace": ["abc"]})) factory = WorkflowFactory(http_request_handler=handler) workflow = factory.create_workflow_from_definition(_yaml(_action(response_headers="Local.H"))) with pytest.raises(DeclarativeActionError): await workflow.run({}) - decl = workflow._runner.state.get(DECLARATIVE_STATE_KEY) - assert decl["Local"]["H"] == {"X-Trace": "abc"} + # After the error and discard(), headers should NOT be in committed state + committed_state = workflow._runner.state.export_state() + # The DECLARATIVE_STATE_KEY may not exist at all if nothing was committed + decl_committed = committed_state.get(DECLARATIVE_STATE_KEY, {}) + # Headers should not be durably persisted after the failed action + assert "Local" not in decl_committed or "H" not in decl_committed.get("Local", {}) # ---------- ConversationId append ------------------------------------------- diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index b3dcebbbc32..a0888c11e05 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -1797,6 +1797,7 @@ async def test_readiness(self) -> None: # region Non-streaming +@pytest.mark.xdist_group("session_isolation") class TestNonStreaming: """ Non-streaming here means that the client requested a non-streaming response, instead of @@ -4086,6 +4087,7 @@ def run_dispatch(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdat return agent +@pytest.mark.xdist_group("session_isolation") class TestMultiTurnMixedContent: """End-to-end multi-turn tests with mixed text and non-text content types.""" @@ -5059,6 +5061,7 @@ async def test_input_item_mcp_approval_response_resolves_to_approval_response(se assert c.approved is False +@pytest.mark.xdist_group("session_isolation") class TestFunctionApprovalRoundTrip: """End-to-end round-trip tests for the function approval flow. @@ -5509,6 +5512,7 @@ async def test_failed_entry_does_not_cache_stack(self) -> None: assert agent.__aenter__.await_count == 2 +@pytest.mark.xdist_group("session_isolation") class TestOAuthConsentSurfacing: async def test_explicit_none_origin_allowlist_accepts_any_safe_https_consent(self) -> None: agent = _make_agent( diff --git a/python/packages/typesafe/agent_framework_typesafe/_tool_calls.py b/python/packages/typesafe/agent_framework_typesafe/_tool_calls.py index cf5508533af..e5b6f4d12e5 100644 --- a/python/packages/typesafe/agent_framework_typesafe/_tool_calls.py +++ b/python/packages/typesafe/agent_framework_typesafe/_tool_calls.py @@ -625,9 +625,7 @@ def _reject_unsupported_schema_constraints( ) -> None: unsupported_constraints = sorted(schema.keys() - supported_keys) if unsupported_constraints: - raise _UnsupportedToolSchema( - f"unsupported {location} schema constraints: {', '.join(unsupported_constraints)}" - ) + raise _UnsupportedToolSchema(f"unsupported {location} schema constraints: {', '.join(unsupported_constraints)}") def _describe_value(value: Any) -> str: