From fbb01312de87a32a7e8632615e5cf97d653bf8c4 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 12:18:31 +0900 Subject: [PATCH 1/7] Harden functional workflow continuation authority Use a versioned opaque single-use token on WorkflowRunResult, validate it before request correlation, consume it immediately before replayed user code, and rotate it on each pause. Carry the same explicit authority through streaming and non-streaming FunctionalWorkflowAgent responses. Files changed: functional workflow/runtime result APIs, functional HITL regression tests, core agent guidance, and the functional HITL sample. Next iteration: enforce pending-state overlap and token-authorized abandonment, then document and test checkpoint authorization boundaries. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 4 +- .../agent_framework/_workflows/_functional.py | 106 +++++++- .../agent_framework/_workflows/_workflow.py | 14 +- .../workflow/test_functional_workflow.py | 242 ++++++++++++++++-- .../03-workflows/functional/hitl_review.py | 8 +- 5 files changed, 334 insertions(+), 40 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 8e633a9334d..4d4c7a652aa 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -199,7 +199,9 @@ agent_framework/ explicit all-output behavior and `intermediate_output_from="all_other"` for visible progress from every output-capable executor not selected by `output_from`. - **`WorkflowRunResult`** - Non-streaming workflow result with Workflow Output `get_outputs()` - and Intermediate Output `get_intermediate_outputs()` accessors + and Intermediate Output `get_intermediate_outputs()` accessors. Functional workflows that pause for + `request_info` also return an opaque, single-use `continuation_token`; pass it with `responses` for an in-memory + response-only resume. - **Orchestrators**: `SequentialOrchestrator`, `ConcurrentOrchestrator`, `GroupChatOrchestrator`, `MagenticOrchestrator`, `HandoffOrchestrator` ## Built-in Providers diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 2ffe99807c0..3ab7a9f2383 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -40,6 +40,7 @@ import hashlib import inspect import logging +import secrets import typing from collections.abc import AsyncIterable, Awaitable, Callable, Sequence from contextvars import ContextVar @@ -48,7 +49,7 @@ from .._feature_stage import ExperimentalFeature, experimental from .._serialization import make_json_safe -from .._types import AgentResponse, AgentResponseUpdate, ResponseStream +from .._types import AgentResponse, AgentResponseUpdate, ContinuationToken, ResponseStream from ..observability import OtelAttr, capture_exception, create_workflow_span from ._checkpoint import CheckpointStorage, WorkflowCheckpoint from ._events import ( @@ -63,6 +64,17 @@ R = TypeVar("R") +_CONTINUATION_KIND: Literal["functional_workflow"] = "functional_workflow" +_CONTINUATION_VERSION: Literal["1"] = "1" +_INVALID_CONTINUATION_AUTHORITY = "Invalid functional workflow continuation authority." + + +class _FunctionalWorkflowContinuationToken(ContinuationToken): + kind: Literal["functional_workflow"] + version: Literal["1"] + token: str + + # ContextVar holding the active RunContext during workflow execution. # ContextVar is per-asyncio-Task, so concurrent workflows each get their own context. _active_run_ctx: ContextVar[RunContext | None] = ContextVar("_active_run_ctx", default=None) @@ -205,8 +217,10 @@ async def request_info( ``ResponseStream`` when ``stream=True``) whose :meth:`~WorkflowRunResult.get_request_info_events` contains the pending request. When the workflow is resumed with - ``run(responses={request_id: value})``, the same function re-executes - and ``request_info`` returns the provided *value* directly. + ``run(responses={request_id: value}, + continuation_token=prior_result.continuation_token)``, the same + function re-executes and ``request_info`` returns the provided *value* + directly. Args: request_data: Arbitrary payload describing what information is @@ -687,6 +701,7 @@ def __init__( self._last_step_cache: dict[tuple[str, int], Any] = {} self._last_step_cache_auto_request_info_counts: dict[tuple[str, int], int] = {} self._last_pending_request_ids: set[str] = set() + self._continuation_nonce: str | None = None # Signature arity is validated once at decoration time. self._non_ctx_param_names = self._classify_signature(func) @@ -740,6 +755,7 @@ def run( *, stream: Literal[True], responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, **kwargs: Any, @@ -752,6 +768,7 @@ def run( *, stream: Literal[False] = ..., responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, @@ -764,6 +781,7 @@ def run( *, stream: bool = False, responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, @@ -787,6 +805,9 @@ def run( responses: HITL responses keyed by ``request_id``, used to resume a workflow that was suspended by :meth:`RunContext.request_info`. + continuation_token: Opaque token returned by the immediately + preceding response-only in-memory run. Required when + *responses* are provided without *checkpoint_id*. checkpoint_id: Identifier of a checkpoint to restore from. Requires *checkpoint_storage* to be set (here or on the decorator). @@ -811,6 +832,9 @@ def run( execution is not allowed). """ self._validate_run_params(message, responses, checkpoint_id) + continuation_nonce: str | None = None + if responses is not None and checkpoint_id is None: + continuation_nonce = self._validate_continuation_authority(continuation_token) # Warn (but don't block) when a fresh message or a checkpoint restore begins while a prior # run left request_info events pending. Mirrors Workflow.run. Delivering responses is the # normal way to complete the pending cycle and is intentionally not warned. @@ -829,7 +853,7 @@ def run( else "those pending requests will be overwritten by the checkpoint's state" ), ) - if responses and checkpoint_id is None: + if responses is not None and checkpoint_id is None: # Require at least one response key to match a currently-pending # request; prevents silent replay against stale state while still # allowing callers to accumulate prior answers across multi-round @@ -848,17 +872,24 @@ def run( f"Provide a response keyed by one of the pending request_ids." ) self._ensure_not_running() + result_continuation_token: list[ContinuationToken | None] = [None] response_stream: ResponseStream[WorkflowEvent[Any], WorkflowRunResult] = ResponseStream( self._run_core( message=message, responses=responses, + continuation_nonce=continuation_nonce, + result_continuation_token=result_continuation_token, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, streaming=stream, **kwargs, ), - finalizer=functools.partial(self._finalize_events, include_status_events=include_status_events), + finalizer=functools.partial( + self._finalize_events, + include_status_events=include_status_events, + continuation_token=result_continuation_token, + ), cleanup_hooks=[self._run_cleanup], ) @@ -917,6 +948,8 @@ async def _run_core( message: Any | None = None, *, responses: dict[str, Any] | None = None, + continuation_nonce: str | None = None, + result_continuation_token: list[ContinuationToken | None], checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, streaming: bool = False, @@ -996,6 +1029,9 @@ async def _on_step_completed() -> None: with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IN_PROGRESS) + if continuation_nonce is not None: + self._consume_continuation_authority(continuation_nonce) + # Execute the user function return_value = await self._execute(ctx, message) @@ -1029,6 +1065,8 @@ async def _on_step_completed() -> None: # Final status if saw_request: self._last_pending_request_ids = set(ctx._pending_requests) + self._rotate_continuation_authority() + result_continuation_token[0] = self._get_continuation_token() with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS) else: @@ -1037,6 +1075,7 @@ async def _on_step_completed() -> None: self._last_step_cache = {} self._last_step_cache_auto_request_info_counts = {} self._last_pending_request_ids = set() + self._continuation_nonce = None with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE) @@ -1047,6 +1086,8 @@ async def _on_step_completed() -> None: self._last_step_cache = dict(ctx._step_cache) self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts) self._last_pending_request_ids = set(ctx._pending_requests) + self._rotate_continuation_authority() + result_continuation_token[0] = self._get_continuation_token() # HITL interruption — yield events collected so far for event in ctx._get_events(): @@ -1209,6 +1250,7 @@ def _finalize_events( events: Sequence[WorkflowEvent[Any]], *, include_status_events: bool = False, + continuation_token: list[ContinuationToken | None], ) -> WorkflowRunResult: filtered: list[WorkflowEvent[Any]] = [] status_events: list[WorkflowEvent[Any]] = [] @@ -1223,7 +1265,39 @@ def _finalize_events( continue filtered.append(ev) - return WorkflowRunResult(filtered, status_events) + return WorkflowRunResult(filtered, status_events, continuation_token[0]) + + def _get_continuation_token(self) -> ContinuationToken | None: + if self._continuation_nonce is None: + return None + return _FunctionalWorkflowContinuationToken( + kind=_CONTINUATION_KIND, + version=_CONTINUATION_VERSION, + token=self._continuation_nonce, + ) + + def _rotate_continuation_authority(self) -> None: + self._continuation_nonce = secrets.token_urlsafe(32) + + def _validate_continuation_authority(self, continuation_token: ContinuationToken | None) -> str: + if ( + self._continuation_nonce is None + or not isinstance(continuation_token, dict) + or set(continuation_token) != {"kind", "version", "token"} + or continuation_token.get("kind") != _CONTINUATION_KIND + or continuation_token.get("version") != _CONTINUATION_VERSION + ): + raise ValueError(_INVALID_CONTINUATION_AUTHORITY) + + token = continuation_token.get("token") + if not isinstance(token, str) or not secrets.compare_digest(token, self._continuation_nonce): + raise ValueError(_INVALID_CONTINUATION_AUTHORITY) + return token + + def _consume_continuation_authority(self, continuation_nonce: str) -> None: + if self._continuation_nonce is None or not secrets.compare_digest(continuation_nonce, self._continuation_nonce): + raise ValueError(_INVALID_CONTINUATION_AUTHORITY) + self._continuation_nonce = None @staticmethod def _validate_run_params( @@ -1341,7 +1415,9 @@ class FunctionalWorkflowAgent: ``request_info`` events emitted by the underlying workflow are surfaced as :class:`FunctionApprovalRequestContent` items (mirroring the graph :class:`WorkflowAgent`), so HITL workflows are callable via this - adapter. Callers resume via ``responses=`` / ``checkpoint_id=``. + adapter. Response-only callers resume via ``responses=`` and the prior + response's ``continuation_token``; checkpoint restores use + ``checkpoint_id=``. Args: workflow: The :class:`FunctionalWorkflow` to wrap. @@ -1386,6 +1462,7 @@ def run( *, stream: Literal[True], responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, **kwargs: Any, @@ -1398,6 +1475,7 @@ def run( *, stream: Literal[False] = ..., responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, **kwargs: Any, @@ -1409,6 +1487,7 @@ def run( *, stream: bool = False, responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, **kwargs: Any, @@ -1423,6 +1502,8 @@ def run( :class:`AgentResponseUpdate` items. responses: HITL responses keyed by ``request_id``, forwarded to the underlying workflow so HITL resumes work via this agent. + continuation_token: Opaque continuation token returned by the + preceding agent response. checkpoint_id: Optional checkpoint to restore from. checkpoint_storage: Override the workflow's default :class:`CheckpointStorage` for this run. @@ -1436,6 +1517,7 @@ def run( return self._run_streaming( messages, responses=responses, + continuation_token=continuation_token, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, **kwargs, @@ -1443,6 +1525,7 @@ def run( return self._run_non_streaming( messages, responses=responses, + continuation_token=continuation_token, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, **kwargs, @@ -1453,6 +1536,7 @@ async def _run_non_streaming( messages: Any | None, *, responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, **kwargs: Any, @@ -1460,6 +1544,7 @@ async def _run_non_streaming( result = await self._workflow.run( messages, responses=responses, + continuation_token=continuation_token, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, **kwargs, @@ -1471,6 +1556,7 @@ def _run_streaming( messages: Any | None, *, responses: dict[str, Any] | None = None, + continuation_token: ContinuationToken | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, **kwargs: Any, @@ -1484,6 +1570,7 @@ def _run_streaming( messages, stream=True, responses=responses, + continuation_token=continuation_token, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, **kwargs, @@ -1513,6 +1600,9 @@ async def _generate_updates() -> AsyncIterable[AgentResponseUpdate]: role="assistant", author_name=agent_name, ) + workflow_result = await workflow_stream.get_final_response() + if workflow_result.continuation_token is not None: + yield AgentResponseUpdate(continuation_token=workflow_result.continuation_token) return ResponseStream( _generate_updates(), @@ -1568,4 +1658,4 @@ def _result_to_agent_response(self, result: WorkflowRunResult) -> AgentResponse: if approval_contents: messages.append(Msg("assistant", approval_contents)) - return AgentResponse(messages=messages) + return AgentResponse(messages=messages, continuation_token=result.continuation_token) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 77060b93a91..dec9258028e 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -17,7 +17,7 @@ from typing import TYPE_CHECKING, Any, Literal, overload from .._sessions import ContextProvider -from .._types import ResponseStream +from .._types import ContinuationToken, ResponseStream from ..exceptions import WorkflowException from ..observability import OtelAttr, capture_exception, create_workflow_span from ._checkpoint import CheckpointStorage @@ -118,11 +118,21 @@ class WorkflowRunResult(list[WorkflowEvent]): - get_request_info_events(): Retrieve external input requests made during execution - get_final_state(): Get the final workflow state (IDLE, IDLE_WITH_PENDING_REQUESTS, etc.) - status_timeline(): Access the complete status event history + + Functional workflows set ``continuation_token`` when execution pauses for + external input. Callers must treat it as opaque and pass it back for the + next response-only in-memory resume. """ - def __init__(self, events: list[WorkflowEvent[Any]], status_events: list[WorkflowEvent[Any]] | None = None) -> None: + def __init__( + self, + events: list[WorkflowEvent[Any]], + status_events: list[WorkflowEvent[Any]] | None = None, + continuation_token: ContinuationToken | None = None, + ) -> None: super().__init__(events) self._status_events: list[WorkflowEvent[Any]] = status_events or [] + self.continuation_token = continuation_token def get_outputs(self) -> list[Any]: """Get all outputs from the workflow run result. diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index c4313e4f4e9..f5d87cad773 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -228,6 +228,60 @@ async def par_wf(x: int) -> tuple[int, int]: class TestHITL: + async def test_response_only_resume_requires_returned_continuation_token(self): + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + feedback = await ctx.request_info({"draft": doc}, response_type=str, request_id="predictable") + return f"Final: {feedback}" + + paused = await review_wf.run("caller data") + + assert paused.continuation_token is not None + assert json.loads(json.dumps(paused.continuation_token)) == paused.continuation_token + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await review_wf.run(responses={"predictable": "stolen"}) + + completed = await review_wf.run( + responses={"predictable": "approved"}, + continuation_token=paused.continuation_token, + ) + + assert completed.get_outputs() == ["Final: approved"] + assert completed.continuation_token is None + + async def test_case_119969_rejects_cross_caller_resume_before_response_correlation(self): + @workflow + async def private_wf(message: str, ctx: RunContext) -> str: + answer = await ctx.request_info( + {"private": message}, + response_type=str, + request_id="private-request-id", + ) + return f"{message}:{answer}" + + paused = await private_wf.run("caller-secret") + assert paused.continuation_token is not None + wrong_token = json.loads(json.dumps(paused.continuation_token)) + wrong_token["token"] = "wrong-token" + + for invalid_token in (None, json.loads("{}"), json.loads('"malformed"'), wrong_token): + with pytest.raises(ValueError) as exc_info: + await private_wf.run( + responses={"guessed-request-id": "attacker-input"}, + continuation_token=invalid_token, + ) + + assert str(exc_info.value) == "Invalid functional workflow continuation authority." + assert "caller-secret" not in str(exc_info.value) + assert "private-request-id" not in str(exc_info.value) + assert "wrong-token" not in str(exc_info.value) + + completed = await private_wf.run( + responses={"private-request-id": "authorized"}, + continuation_token=paused.continuation_token, + ) + assert completed.get_outputs() == ["caller-secret:authorized"] + async def test_request_info_interrupts(self): @workflow async def review_wf(doc: str, ctx: RunContext) -> str: @@ -252,7 +306,10 @@ async def review_wf(doc: str, ctx: RunContext) -> str: assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS # Phase 2: resume with response - result2 = await review_wf.run(responses={"req1": "Looks great!"}) + result2 = await review_wf.run( + responses={"req1": "Looks great!"}, + continuation_token=result1.continuation_token, + ) outputs = result2.get_outputs() assert outputs == ["Final: Looks great!"] assert result2.get_final_state() == WorkflowRunState.IDLE @@ -289,7 +346,10 @@ async def review_wf(doc: str, ctx: RunContext) -> str: caplog.clear() with caplog.at_level(logging.WARNING): - result2 = await review_wf.run(responses={"req1": "Looks great!"}) + result2 = await review_wf.run( + responses={"req1": "Looks great!"}, + continuation_token=result1.continuation_token, + ) assert result2.get_final_state() == WorkflowRunState.IDLE assert "still pending" not in caplog.text @@ -305,7 +365,10 @@ async def review_wf(doc: str, ctx) -> str: # pyright: ignore[reportMissingParam result1 = await review_wf.run("my doc") assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS - result2 = await review_wf.run(responses={"req1": "LGTM"}) + result2 = await review_wf.run( + responses={"req1": "LGTM"}, + continuation_token=result1.continuation_token, + ) assert result2.get_outputs() == ["Final: LGTM"] async def test_multiple_sequential_interrupts(self): @@ -319,15 +382,57 @@ async def multi_hitl(data: str, ctx: RunContext) -> str: result1 = await multi_hitl.run("start") assert len(result1.get_request_info_events()) == 1 assert result1.get_request_info_events()[0].request_id == "r1" + assert result1.continuation_token is not None # Phase 2: respond to first, hits second - result2 = await multi_hitl.run(responses={"r1": "A"}) + result2 = await multi_hitl.run( + responses={"r1": "A"}, + continuation_token=result1.continuation_token, + ) assert len(result2.get_request_info_events()) == 1 assert result2.get_request_info_events()[0].request_id == "r2" + assert result2.continuation_token is not None + assert result2.continuation_token != result1.continuation_token + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await multi_hitl.run( + responses={"r1": "A", "r2": "stale"}, + continuation_token=result1.continuation_token, + ) # Phase 3: respond to second - result3 = await multi_hitl.run(responses={"r1": "A", "r2": "B"}) + result3 = await multi_hitl.run( + responses={"r1": "A", "r2": "B"}, + continuation_token=result2.continuation_token, + ) assert result3.get_outputs() == ["A+B"] + assert result3.continuation_token is None + + async def test_continuation_token_is_consumed_before_resumed_user_code_fails(self): + resumed_user_code_started = False + + @workflow + async def failing_resume(data: str, ctx: RunContext) -> str: + nonlocal resumed_user_code_started + answer = await ctx.request_info(data, response_type=str, request_id="r1") + resumed_user_code_started = True + raise RuntimeError(f"resume failed after {answer}") + + paused = await failing_resume.run("input") + assert paused.continuation_token is not None + + with pytest.raises(RuntimeError, match="resume failed after response"): + await failing_resume.run( + responses={"r1": "response"}, + continuation_token=paused.continuation_token, + ) + assert resumed_user_code_started + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await failing_resume.run( + responses={"r1": "response"}, + continuation_token=paused.continuation_token, + ) async def test_request_info_auto_generates_id(self): @workflow @@ -460,6 +565,25 @@ async def wf(x: int, ctx: RunContext) -> int: await wf.run(1) assert streaming_flag is False + async def test_streaming_final_response_carries_continuation_token(self): + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + feedback = await ctx.request_info(doc, response_type=str, request_id="r1") + return f"{doc}:{feedback}" + + paused_stream = review_wf.run("draft", stream=True) + paused = await paused_stream.get_final_response() + assert paused.continuation_token is not None + + completed_stream = review_wf.run( + responses={"r1": "approved"}, + continuation_token=paused.continuation_token, + stream=True, + ) + completed = await completed_stream.get_final_response() + assert completed.get_outputs() == ["draft:approved"] + assert completed.continuation_token is None + # --------------------------------------------------------------------------- # Step passthrough outside workflow @@ -1180,7 +1304,10 @@ async def wf(doc: str) -> str: assert result1.get_request_info_events()[0].request_id == "s1" # Phase 2: resume - result2 = await wf.run(responses={"s1": "LGTM"}) + result2 = await wf.run( + responses={"s1": "LGTM"}, + continuation_token=result1.continuation_token, + ) assert result2.get_outputs() == ["reviewed: LGTM"] async def test_step_works_outside_workflow_with_explicit_ctx(self): @@ -1255,11 +1382,14 @@ async def wf(doc: str, ctx: RunContext) -> str: return f"got: {val}" # Phase 1 - await wf.run("start") + paused = await wf.run("start") # Phase 2: resume with None response — should warn but still work with caplog_context(logging.getLogger("agent_framework._workflows._functional")) as logs: - result = await wf.run(responses={"r1": None}) + result = await wf.run( + responses={"r1": None}, + continuation_token=paused.continuation_token, + ) assert result.get_outputs() == ["got: None"] assert any("None" in msg and "r1" in msg for msg in logs) @@ -1272,8 +1402,11 @@ async def wf(x: int, ctx: RunContext) -> str: val = await ctx.request_info("need data", response_type=str, request_id="r1") return f"value={val}" - await wf.run(1) - result = await wf.run(responses={"r1": None}) + paused = await wf.run(1) + result = await wf.run( + responses={"r1": None}, + continuation_token=paused.continuation_token, + ) assert result.get_outputs() == ["value=None"] @@ -1333,7 +1466,10 @@ async def wf(x: int) -> str: assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS # Phase 2: resume — step_a should be bypassed, step_b re-executes - result2 = await wf.run(responses={"r1": "ok"}) + result2 = await wf.run( + responses={"r1": "ok"}, + continuation_token=result1.continuation_token, + ) assert call_count_a == 1 # step_a not called again assert result2.get_outputs() == ["6:ok"] @@ -1363,7 +1499,10 @@ async def wf(x: int) -> str: assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS # Phase 2: resume - result2 = await wf.run(responses={"rev": "LGTM"}) + result2 = await wf.run( + responses={"rev": "LGTM"}, + continuation_token=result1.continuation_token, + ) assert result2.get_outputs() == ["reviewed(30):LGTM"] # Phase 3: restore from latest checkpoint -- both steps should be bypassed @@ -1388,10 +1527,13 @@ async def needs_feedback(doc: str, ctx: RunContext) -> str: async def wf(doc: str) -> str: return await needs_feedback(doc) - await wf.run("draft") + paused = await wf.run("draft") with caplog_context(logging.getLogger("agent_framework._workflows._functional")) as logs: - result = await wf.run(responses={"r1": None}) + result = await wf.run( + responses={"r1": None}, + continuation_token=paused.continuation_token, + ) assert result.get_outputs() == ["got:None"] assert any("None" in msg and "r1" in msg for msg in logs) @@ -1436,7 +1578,10 @@ async def wf(x: int, ctx: RunContext) -> str: assert rid # non-empty # Resume with the id the caller just received. - result2 = await wf.run(responses={rid: "hello"}) + result2 = await wf.run( + responses={rid: "hello"}, + continuation_token=result1.continuation_token, + ) assert result2.get_final_state() == WorkflowRunState.IDLE assert result2.get_outputs() == ["got:hello"] @@ -1449,10 +1594,16 @@ async def wf(x: int, ctx: RunContext) -> str: r1 = await wf.run(1) rid1 = r1.get_request_info_events()[0].request_id - r2 = await wf.run(responses={rid1: "A"}) + r2 = await wf.run( + responses={rid1: "A"}, + continuation_token=r1.continuation_token, + ) rid2 = r2.get_request_info_events()[0].request_id assert rid1 != rid2 - r3 = await wf.run(responses={rid1: "A", rid2: "B"}) + r3 = await wf.run( + responses={rid1: "A", rid2: "B"}, + continuation_token=r2.continuation_token, + ) assert r3.get_outputs() == ["A/B"] async def test_cached_step_advances_auto_request_id_counter(self): @@ -1478,12 +1629,18 @@ async def wf(value: int) -> str: first_request_id = first_run.get_request_info_events()[0].request_id assert first_request_id == "auto::0" - second_run = await wf.run(responses={first_request_id: "A"}) + second_run = await wf.run( + responses={first_request_id: "A"}, + continuation_token=first_run.continuation_token, + ) second_request_id = second_run.get_request_info_events()[0].request_id assert second_request_id == "auto::1" completed_call_count = call_count - final_run = await wf.run(responses={first_request_id: "A", second_request_id: "B"}) + final_run = await wf.run( + responses={first_request_id: "A", second_request_id: "B"}, + continuation_token=second_run.continuation_token, + ) assert call_count == completed_call_count assert final_run.get_outputs() == ["A/B"] @@ -1501,9 +1658,15 @@ async def wf(x: int, ctx: RunContext) -> str: b = await ctx.request_info("q2", response_type=str, request_id="r2") return f"{a}/{b}" - await wf.run(1) - await wf.run(responses={"r1": "A"}) - result = await wf.run(responses={"r1": "A", "r2": "B"}) + first_run = await wf.run(1) + second_run = await wf.run( + responses={"r1": "A"}, + continuation_token=first_run.continuation_token, + ) + result = await wf.run( + responses={"r1": "A", "r2": "B"}, + continuation_token=second_run.continuation_token, + ) assert result.get_final_state() == WorkflowRunState.IDLE # Latest checkpoint must show no pending requests. checkpoints = await storage.list_checkpoints(workflow_name="wf") @@ -1551,7 +1714,7 @@ async def wf(x: int) -> int: return x * 2 await wf.run(5) # clean completion, no pending requests - with pytest.raises(ValueError, match="no pending request_info"): + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): await wf.run(responses={"stale": "x"}) async def test_responses_mismatched_key_raises(self): @@ -1559,9 +1722,12 @@ async def test_responses_mismatched_key_raises(self): async def wf(x: int, ctx: RunContext) -> str: return await ctx.request_info("q", response_type=str, request_id="r1") - await wf.run(1) # interrupts with r1 pending + paused = await wf.run(1) # interrupts with r1 pending with pytest.raises(ValueError, match="do not answer"): - await wf.run(responses={"definitely_not_r1": "x"}) + await wf.run( + responses={"definitely_not_r1": "x"}, + continuation_token=paused.continuation_token, + ) class TestReservedStateKeys: @@ -1719,9 +1885,13 @@ async def wf(x: str, ctx: RunContext) -> str: agent = wf.as_agent() # First phase: suspend - await agent.run("topic") + paused = await agent.run("topic") + assert paused.continuation_token is not None # Second phase: resume via the agent surface - response = await agent.run(responses={"rid-1": "answered"}) + response = await agent.run( + responses={"rid-1": "answered"}, + continuation_token=paused.continuation_token, + ) # Agent's final response should contain the workflow's text output. text_blobs: list[str] = [] for message in response.messages: @@ -1731,6 +1901,24 @@ async def wf(x: str, ctx: RunContext) -> str: text_blobs.append(text) assert any("got:answered" in t for t in text_blobs) + async def test_streaming_resume_carries_continuation_token(self): + @workflow + async def wf(x: str, ctx: RunContext) -> str: + answer = await ctx.request_info(x, response_type=str, request_id="rid-1") + return f"got:{answer}" + + agent = wf.as_agent() + paused = await agent.run("topic", stream=True).get_final_response() + assert paused.continuation_token is not None + + completed = await agent.run( + responses={"rid-1": "answered"}, + continuation_token=paused.continuation_token, + stream=True, + ).get_final_response() + assert completed.text == "got:answered" + assert completed.continuation_token is None + class TestRunDocstringAllowsResponsesAndCheckpoint: """Regression for bug_010: docstring must permit responses+checkpoint_id combo.""" diff --git a/python/samples/03-workflows/functional/hitl_review.py b/python/samples/03-workflows/functional/hitl_review.py index 39f2dae8853..78e473bbf39 100644 --- a/python/samples/03-workflows/functional/hitl_review.py +++ b/python/samples/03-workflows/functional/hitl_review.py @@ -3,7 +3,7 @@ """Human-in-the-loop review pipeline using functional workflows. Demonstrates ctx.request_info() for pausing the workflow to wait for -external input and resuming with run(responses={...}). +external input and resuming with the returned continuation token. HITL works with or without @step. The difference is what happens on resume: - Without @step: every function re-executes from the top (fine for cheap calls). @@ -70,11 +70,15 @@ async def main(): requests = result1.get_request_info_events() print(f"Pending request: {requests[0].request_id}") + assert result1.continuation_token is not None # Phase 2: Resume with the human's response print("\n=== Phase 2: Resume with feedback ===") print("(write_draft should NOT execute again — saved by @step)") - result2 = await review_pipeline.run(responses={"review_request": "Add more details about alignment research"}) + result2 = await review_pipeline.run( + responses={"review_request": "Add more details about alignment research"}, + continuation_token=result1.continuation_token, + ) print(f"State: {result2.get_final_state()}") print(f"Output: {result2.get_outputs()[0]}") From ada453fec70be694b5fa426b10cc7396ec79d166 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 12:29:25 +0900 Subject: [PATCH 2/7] Enforce one pending functional continuation Reject fresh messages and checkpoint restores while an in-memory continuation is pending. Add token-authorized abandonment on FunctionalWorkflow and FunctionalWorkflowAgent, and clear retained replay state atomically when authority is consumed while preserving the active message for token rotation and checkpoints. Files changed: functional workflow runtime and agent adapter, functional lifecycle regression tests, and core workflow guidance. Next iteration: preserve and document authorized checkpoint continuation boundaries. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 5 +- .../agent_framework/_workflows/_functional.py | 84 ++++++----- .../workflow/test_functional_workflow.py | 134 ++++++++++++++++-- 3 files changed, 177 insertions(+), 46 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 4d4c7a652aa..668c5888203 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -201,7 +201,10 @@ agent_framework/ - **`WorkflowRunResult`** - Non-streaming workflow result with Workflow Output `get_outputs()` and Intermediate Output `get_intermediate_outputs()` accessors. Functional workflows that pause for `request_info` also return an opaque, single-use `continuation_token`; pass it with `responses` for an in-memory - response-only resume. + response-only resume. Each `FunctionalWorkflow` instance retains at most one in-memory continuation: callers must + resume it or call `abandon_continuation(token)` before starting new input or restoring a checkpoint. Use separate + workflow instances for independent in-memory runs; `FunctionalWorkflowAgent.abandon_continuation` delegates the + same operation. - **Orchestrators**: `SequentialOrchestrator`, `ConcurrentOrchestrator`, `GroupChatOrchestrator`, `MagenticOrchestrator`, `HandoffOrchestrator` ## Built-in Providers diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 3ab7a9f2383..90d3284349e 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -660,6 +660,11 @@ class FunctionalWorkflow: edge wiring is involved. Native Python control flow (``if``/``else``, ``for``, ``asyncio.gather``) is used for branching and parallelism. + A workflow instance retains at most one in-memory continuation. Resume + or explicitly abandon a pending continuation before starting new input + or restoring a checkpoint on that instance. Use separate workflow + instances for independent in-memory runs. + Args: func: The async function that implements the workflow logic. name: Display name for the workflow. Defaults to ``func.__name__``. @@ -828,30 +833,18 @@ def run( Raises: ValueError: If the combination of *message*, *responses*, and *checkpoint_id* is invalid. - RuntimeError: If the workflow is already running (concurrent - execution is not allowed). + RuntimeError: If the workflow is already running, or if new input + or a checkpoint restore is attempted while an in-memory + continuation is pending. """ self._validate_run_params(message, responses, checkpoint_id) continuation_nonce: str | None = None if responses is not None and checkpoint_id is None: continuation_nonce = self._validate_continuation_authority(continuation_token) - # Warn (but don't block) when a fresh message or a checkpoint restore begins while a prior - # run left request_info events pending. Mirrors Workflow.run. Delivering responses is the - # normal way to complete the pending cycle and is intentionally not warned. if (message is not None or checkpoint_id is not None) and self._last_pending_request_ids: - logger.warning( - "Workflow %s received %s while %d request_info event(s) are still pending from an " - "unfinished request/response cycle; %s. Deliver responses (responses=...) to complete " - "the pending cycle before starting new input.", - self.name, - "a fresh message" if message is not None else "a checkpoint restore", - len(self._last_pending_request_ids), - ( - "those requests remain answerable, but this run advances workflow state, so a " - "response that arrives later may apply to a workflow that has moved on" - if message is not None - else "those pending requests will be overwritten by the checkpoint's state" - ), + raise RuntimeError( + "Cannot start or restore a functional workflow run while an in-memory continuation is pending. " + "Resume or abandon the pending continuation first." ) if responses is not None and checkpoint_id is None: # Require at least one response key to match a currently-pending @@ -939,6 +932,25 @@ def as_agent( **kwargs, ) + def abandon_continuation(self, continuation_token: ContinuationToken | None = None) -> None: + """Abandon the pending in-memory continuation. + + Successful abandonment consumes the token and clears the retained + message, step cache, request metadata, and pending requests. A failed + attempt leaves the continuation unchanged. + + Args: + continuation_token: Opaque token returned by the pending run. + + Raises: + RuntimeError: If the workflow is currently running. + ValueError: If the token does not authorize the current pending continuation. + """ + if self._is_running: + raise RuntimeError("Cannot abandon a continuation while the functional workflow is running.") + continuation_nonce = self._validate_continuation_authority(continuation_token) + self._consume_continuation_authority(continuation_nonce) + # ------------------------------------------------------------------ # Internal execution # ------------------------------------------------------------------ @@ -996,10 +1008,6 @@ async def _run_core( ctx._step_cache = dict(self._last_step_cache) ctx._step_cache_auto_request_info_counts = dict(self._last_step_cache_auto_request_info_counts) - # Store message for future replays - if message is not None: - self._last_message = message - # Set responses for replay if responses: ctx._set_responses(responses) @@ -1010,7 +1018,7 @@ async def _run_core( if storage is not None: async def _on_step_completed() -> None: - ckpt_chain[0] = await self._save_checkpoint(ctx, storage, ckpt_chain[0]) + ckpt_chain[0] = await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) ctx._on_step_completed = _on_step_completed @@ -1060,10 +1068,11 @@ async def _on_step_completed() -> None: # Save final checkpoint if storage is available if storage is not None: - await self._save_checkpoint(ctx, storage, ckpt_chain[0]) + await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) # Final status if saw_request: + self._last_message = message self._last_pending_request_ids = set(ctx._pending_requests) self._rotate_continuation_authority() result_continuation_token[0] = self._get_continuation_token() @@ -1071,11 +1080,7 @@ async def _on_step_completed() -> None: yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS) else: # Clean completion — drop cross-run replay state. - self._last_message = None - self._last_step_cache = {} - self._last_step_cache_auto_request_info_counts = {} - self._last_pending_request_ids = set() - self._continuation_nonce = None + self._clear_continuation_state() with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE) @@ -1083,6 +1088,7 @@ async def _on_step_completed() -> None: except WorkflowInterrupted: # Persist step cache for response-only replay + self._last_message = message self._last_step_cache = dict(ctx._step_cache) self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts) self._last_pending_request_ids = set(ctx._pending_requests) @@ -1100,7 +1106,7 @@ async def _on_step_completed() -> None: # Save checkpoint if storage is not None: - await self._save_checkpoint(ctx, storage, ckpt_chain[0]) + await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS) @@ -1184,12 +1190,13 @@ async def _save_checkpoint( self, ctx: RunContext, storage: CheckpointStorage, + original_message: Any, previous_checkpoint_id: str | None = None, ) -> str: state = dict(ctx._state) state["_step_cache"] = ctx._export_step_cache() state["_step_cache_auto_request_info_counts"] = ctx._export_step_cache_auto_request_info_counts() - state["_original_message"] = self._last_message + state["_original_message"] = original_message checkpoint = WorkflowCheckpoint( workflow_name=self.name, @@ -1297,6 +1304,13 @@ def _validate_continuation_authority(self, continuation_token: ContinuationToken def _consume_continuation_authority(self, continuation_nonce: str) -> None: if self._continuation_nonce is None or not secrets.compare_digest(continuation_nonce, self._continuation_nonce): raise ValueError(_INVALID_CONTINUATION_AUTHORITY) + self._clear_continuation_state() + + def _clear_continuation_state(self) -> None: + self._last_message = None + self._last_step_cache = {} + self._last_step_cache_auto_request_info_counts = {} + self._last_pending_request_ids = set() self._continuation_nonce = None @staticmethod @@ -1417,7 +1431,8 @@ class FunctionalWorkflowAgent: :class:`WorkflowAgent`), so HITL workflows are callable via this adapter. Response-only callers resume via ``responses=`` and the prior response's ``continuation_token``; checkpoint restores use - ``checkpoint_id=``. + ``checkpoint_id=``. :meth:`abandon_continuation` delegates token-authorized + abandonment to the wrapped workflow. Args: workflow: The :class:`FunctionalWorkflow` to wrap. @@ -1455,6 +1470,11 @@ def pending_requests(self) -> dict[str, WorkflowEvent[Any]]: """Pending request_info events emitted during the last run.""" return self._pending_requests + def abandon_continuation(self, continuation_token: ContinuationToken | None = None) -> None: + """Abandon the wrapped workflow's pending in-memory continuation.""" + self._workflow.abandon_continuation(continuation_token) + self._pending_requests = {} + @overload def run( self, diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index f5d87cad773..15a6fbd1ed9 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -314,24 +314,98 @@ async def review_wf(doc: str, ctx: RunContext) -> str: assert outputs == ["Final: Looks great!"] assert result2.get_final_state() == WorkflowRunState.IDLE - async def test_fresh_message_while_pending_requests_warns(self, caplog: pytest.LogCaptureFixture) -> None: - """A fresh message while request_info events are pending is allowed but logs a warning.""" - + async def test_fresh_message_while_pending_requests_is_rejected_without_losing_continuation(self) -> None: @workflow async def review_wf(doc: str, ctx: RunContext) -> str: feedback = await ctx.request_info({"draft": doc}, response_type=str, request_id="req1") return f"Final: {feedback}" - result1 = await review_wf.run("my doc") - assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS - - # Starting fresh input while a request is pending does not abandon it, but advances - # workflow state so a later response may apply to a moved-on workflow -> warn (but proceed). - with caplog.at_level(logging.WARNING): + paused = await review_wf.run("my doc") + with pytest.raises(RuntimeError, match="(?i)resume or abandon the pending continuation"): await review_wf.run("another doc") - assert "request_info event(s) are still pending" in caplog.text - assert "a fresh message" in caplog.text + completed = await review_wf.run( + responses={"req1": "approved"}, + continuation_token=paused.continuation_token, + ) + assert completed.get_outputs() == ["Final: approved"] + + async def test_checkpoint_restore_while_pending_is_rejected_without_losing_continuation(self) -> None: + storage = InMemoryCheckpointStorage() + + @workflow(checkpoint_storage=storage) + async def review_wf(doc: str, ctx: RunContext) -> str: + feedback = await ctx.request_info({"draft": doc}, response_type=str, request_id="req1") + return f"{doc}: {feedback}" + + paused = await review_wf.run("original") + checkpoints = await storage.list_checkpoints(workflow_name="review_wf") + + with pytest.raises(RuntimeError, match="(?i)resume or abandon the pending continuation"): + await review_wf.run(checkpoint_id=checkpoints[0].checkpoint_id) + + completed = await review_wf.run( + responses={"req1": "approved"}, + continuation_token=paused.continuation_token, + ) + assert completed.get_outputs() == ["original: approved"] + + async def test_abandon_continuation_requires_current_token_and_preserves_state_on_failure(self) -> None: + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + feedback = await ctx.request_info(doc, response_type=str, request_id="req1") + return f"{doc}: {feedback}" + + paused = await review_wf.run("original") + assert paused.continuation_token is not None + wrong_token = json.loads(json.dumps(paused.continuation_token)) + wrong_token["token"] = "wrong" + + for invalid_token in (None, json.loads("{}"), json.loads('"malformed"'), wrong_token): + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + review_wf.abandon_continuation(invalid_token) + + completed = await review_wf.run( + responses={"req1": "approved"}, + continuation_token=paused.continuation_token, + ) + assert completed.get_outputs() == ["original: approved"] + + async def test_abandon_continuation_clears_replay_state_and_allows_fresh_run(self) -> None: + step_calls = 0 + + @step + async def prepare(doc: str) -> str: + nonlocal step_calls + step_calls += 1 + return f"prepared:{doc}" + + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + prepared = await prepare(doc) + feedback = await ctx.request_info(prepared, response_type=str) + return f"{prepared}: {feedback}" + + abandoned = await review_wf.run("original") + assert abandoned.continuation_token is not None + assert abandoned.get_request_info_events()[0].request_id == "auto::0" + assert step_calls == 1 + + review_wf.abandon_continuation(abandoned.continuation_token) + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + review_wf.abandon_continuation(abandoned.continuation_token) + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await review_wf.run( + responses={"auto::0": "stale"}, + continuation_token=abandoned.continuation_token, + ) + + fresh = await review_wf.run("new") + fresh_request = fresh.get_request_info_events()[0] + assert fresh_request.request_id == "auto::0" + assert fresh_request.data == "prepared:new" + assert step_calls == 2 async def test_responses_while_pending_requests_does_not_warn(self, caplog: pytest.LogCaptureFixture) -> None: """Delivering responses is the normal completion path and must not warn.""" @@ -376,7 +450,7 @@ async def test_multiple_sequential_interrupts(self): async def multi_hitl(data: str, ctx: RunContext) -> str: r1 = await ctx.request_info("step1", response_type=str, request_id="r1") r2 = await ctx.request_info("step2", response_type=str, request_id="r2") - return f"{r1}+{r2}" + return f"{data}:{r1}+{r2}" # Phase 1: first interrupt result1 = await multi_hitl.run("start") @@ -405,7 +479,7 @@ async def multi_hitl(data: str, ctx: RunContext) -> str: responses={"r1": "A", "r2": "B"}, continuation_token=result2.continuation_token, ) - assert result3.get_outputs() == ["A+B"] + assert result3.get_outputs() == ["start:A+B"] assert result3.continuation_token is None async def test_continuation_token_is_consumed_before_resumed_user_code_fails(self): @@ -434,6 +508,10 @@ async def failing_resume(data: str, ctx: RunContext) -> str: continuation_token=paused.continuation_token, ) + fresh = await failing_resume.run("fresh") + assert fresh.continuation_token is not None + assert fresh.get_request_info_events()[0].data == "fresh" + async def test_request_info_auto_generates_id(self): @workflow async def auto_id_wf(x: int, ctx: RunContext) -> None: @@ -712,6 +790,7 @@ async def hitl_wf(doc: str, ctx: RunContext) -> str: # Phase 1: interrupt result1 = await hitl_wf.run("draft text") assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + hitl_wf.abandon_continuation(result1.continuation_token) # Get checkpoint checkpoints = await storage.list_checkpoints(workflow_name="hitl_wf") @@ -742,6 +821,7 @@ async def stateful_wf(x: int, ctx: RunContext) -> str: # Phase 1 result1 = await stateful_wf.run(1) assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + stateful_wf.abandon_continuation(result1.continuation_token) # Phase 2: restore and respond checkpoints = await storage.list_checkpoints(workflow_name="stateful_wf") @@ -1919,6 +1999,34 @@ async def wf(x: str, ctx: RunContext) -> str: assert completed.text == "got:answered" assert completed.continuation_token is None + async def test_agent_can_abandon_pending_continuation(self) -> None: + @workflow + async def wf(x: str, ctx: RunContext) -> str: + answer = await ctx.request_info(x, response_type=str, request_id="rid-1") + return f"{x}:{answer}" + + agent = wf.as_agent() + abandoned = await agent.run("original") + assert abandoned.continuation_token is not None + + wrong_token = json.loads(json.dumps(abandoned.continuation_token)) + wrong_token["token"] = "wrong" + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + agent.abandon_continuation(wrong_token) + assert "rid-1" in agent.pending_requests + + agent.abandon_continuation(abandoned.continuation_token) + assert agent.pending_requests == {} + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await agent.run( + responses={"rid-1": "stale"}, + continuation_token=abandoned.continuation_token, + ) + fresh = await agent.run("new") + assert fresh.continuation_token is not None + assert fresh.continuation_token != abandoned.continuation_token + class TestRunDocstringAllowsResponsesAndCheckpoint: """Regression for bug_010: docstring must permit responses+checkpoint_id combo.""" From a602ea9310fc0c202bfabdf9edc4f68988f7be87 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 12:33:56 +0900 Subject: [PATCH 3/7] Preserve authorized functional checkpoint continuation Treat checkpoint restore as a host- and storage-authorized path independent of process-local continuation tokens, and issue fresh authority whenever restored execution pauses again. Cover default and per-run storage, deterministic and custom request IDs, token rotation, and checkpoint-plus-response restore. Files changed: functional workflow and checkpoint interface guidance, functional checkpoint lifecycle tests, the functional HITL sample, and core workflow guidance. Next iteration: run the final repository-wide Python validation gates. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 5 +- .../agent_framework/_workflows/_checkpoint.py | 8 ++- .../agent_framework/_workflows/_functional.py | 30 +++++++-- .../workflow/test_functional_workflow.py | 62 +++++++++++++++++++ .../03-workflows/functional/hitl_review.py | 5 +- 5 files changed, 101 insertions(+), 9 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 668c5888203..611cb499f0d 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -204,7 +204,10 @@ agent_framework/ response-only resume. Each `FunctionalWorkflow` instance retains at most one in-memory continuation: callers must resume it or call `abandon_continuation(token)` before starting new input or restoring a checkpoint. Use separate workflow instances for independent in-memory runs; `FunctionalWorkflowAgent.abandon_continuation` delegates the - same operation. + same operation. Checkpoint restoration is a separate host-authorized path: checkpoint IDs only locate persisted + state, while the host or storage adapter owns authorization and tenant isolation. A restored functional workflow + that pauses returns fresh process-local continuation authority. This does not alter graph-workflow request-info + authoritative resolution from PR #7500. - **Orchestrators**: `SequentialOrchestrator`, `ConcurrentOrchestrator`, `GroupChatOrchestrator`, `MagenticOrchestrator`, `HandoffOrchestrator` ## Built-in Providers diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 2b267979e99..3dc600cb915 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -127,7 +127,13 @@ def from_dict(cls, data: Mapping[str, Any]) -> WorkflowCheckpoint: class CheckpointStorage(Protocol): - """Protocol for checkpoint storage backends.""" + """Protocol for checkpoint storage backends. + + Checkpoint IDs locate persisted workflow state; they are not authentication + or authorization credentials, even when represented as UUIDs. Hosts and + storage adapters are responsible for authorizing checkpoint operations and + isolating checkpoint data between tenants. + """ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: """Save a checkpoint and return its ID. diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 90d3284349e..263f2bdeea1 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -665,6 +665,17 @@ class FunctionalWorkflow: or restoring a checkpoint on that instance. Use separate workflow instances for independent in-memory runs. + Checkpoint restoration is a distinct, host-authorized continuation path + and does not require continuation authority from the process that created + the checkpoint. The host or checkpoint-storage adapter must authorize + access and enforce tenant isolation; a checkpoint ID, including a UUID, is + only a locator. If a restored run pauses, its result carries fresh + process-local continuation authority for later response-only runs. + + These continuation rules apply only to functional workflows. Graph + workflow request-info authoritative resolution remains a separate concern, + as hardened in PR #7500. + Args: func: The async function that implements the workflow logic. name: Display name for the workflow. Defaults to ``func.__name__``. @@ -811,11 +822,14 @@ def run( resume a workflow that was suspended by :meth:`RunContext.request_info`. continuation_token: Opaque token returned by the immediately - preceding response-only in-memory run. Required when - *responses* are provided without *checkpoint_id*. + preceding result for the pending in-memory continuation. + Required when *responses* are provided without *checkpoint_id*. checkpoint_id: Identifier of a checkpoint to restore from. Requires *checkpoint_storage* to be set (here or on the - decorator). + decorator). Checkpoint restoration does not use prior + process-local continuation authority; the host or storage + adapter is responsible for authorizing checkpoint access and + tenant isolation. checkpoint_storage: Override the default checkpoint storage for this run. include_status_events: When ``True`` (non-streaming only), @@ -1431,8 +1445,10 @@ class FunctionalWorkflowAgent: :class:`WorkflowAgent`), so HITL workflows are callable via this adapter. Response-only callers resume via ``responses=`` and the prior response's ``continuation_token``; checkpoint restores use - ``checkpoint_id=``. :meth:`abandon_continuation` delegates token-authorized - abandonment to the wrapped workflow. + ``checkpoint_id=`` after the host or storage adapter authorizes access. + A restored run that pauses returns fresh process-local authority through + the agent response's ``continuation_token``. :meth:`abandon_continuation` + delegates token-authorized abandonment to the wrapped workflow. Args: workflow: The :class:`FunctionalWorkflow` to wrap. @@ -1524,7 +1540,9 @@ def run( the underlying workflow so HITL resumes work via this agent. continuation_token: Opaque continuation token returned by the preceding agent response. - checkpoint_id: Optional checkpoint to restore from. + checkpoint_id: Optional host-authorized checkpoint to restore + from. A checkpoint ID locates state; it is not an + authorization credential. checkpoint_storage: Override the workflow's default :class:`CheckpointStorage` for this run. **kwargs: Extra keyword arguments forwarded to the workflow run. diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 15a6fbd1ed9..c299c027002 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -716,6 +716,68 @@ async def wf(x: int, ctx: RunContext) -> str: class TestCheckpointing: + async def test_restored_checkpoint_issues_process_local_continuation_authority(self): + storage = InMemoryCheckpointStorage() + + async def review(doc: str, ctx: RunContext) -> str: + first = await ctx.request_info({"draft": doc}, response_type=str) + second = await ctx.request_info({"first": first}, response_type=str, request_id="final-review") + return f"{doc}:{first}:{second}" + + original_process = workflow(checkpoint_storage=storage)(review) + original_pause = await original_process.run("draft") + checkpoint = (await storage.list_checkpoints(workflow_name="review"))[-1] + + restored_process = workflow(checkpoint_storage=storage)(review) + restored_pause = await restored_process.run(checkpoint_id=checkpoint.checkpoint_id) + + assert restored_pause.get_request_info_events()[0].request_id == "auto::0" + assert restored_pause.continuation_token is not None + assert restored_pause.continuation_token != original_pause.continuation_token + + for invalid_token in (None, original_pause.continuation_token): + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await restored_process.run( + responses={"auto::0": "approved"}, + continuation_token=invalid_token, + ) + + second_pause = await restored_process.run( + responses={"auto::0": "approved"}, + continuation_token=restored_pause.continuation_token, + ) + assert second_pause.get_request_info_events()[0].request_id == "final-review" + assert second_pause.continuation_token is not None + assert second_pause.continuation_token != restored_pause.continuation_token + + completed = await restored_process.run( + responses={"auto::0": "approved", "final-review": "ship it"}, + continuation_token=second_pause.continuation_token, + ) + assert completed.get_outputs() == ["draft:approved:ship it"] + assert completed.continuation_token is None + + async def test_runtime_storage_override_restores_checkpoint_with_responses_without_token(self): + storage = InMemoryCheckpointStorage() + + async def review(doc: str, ctx: RunContext) -> str: + feedback = await ctx.request_info(doc, response_type=str, request_id="predictable-review") + return f"{doc}:{feedback}" + + original_process = workflow(review) + await original_process.run("draft", checkpoint_storage=storage) + checkpoint = (await storage.list_checkpoints(workflow_name="review"))[-1] + + restored_process = workflow(review) + completed = await restored_process.run( + checkpoint_id=checkpoint.checkpoint_id, + responses={"predictable-review": "approved"}, + checkpoint_storage=storage, + ) + + assert completed.get_outputs() == ["draft:approved"] + assert completed.continuation_token is None + async def test_checkpoint_save_and_restore(self): storage = InMemoryCheckpointStorage() diff --git a/python/samples/03-workflows/functional/hitl_review.py b/python/samples/03-workflows/functional/hitl_review.py index 78e473bbf39..bd360854924 100644 --- a/python/samples/03-workflows/functional/hitl_review.py +++ b/python/samples/03-workflows/functional/hitl_review.py @@ -72,7 +72,10 @@ async def main(): print(f"Pending request: {requests[0].request_id}") assert result1.continuation_token is not None - # Phase 2: Resume with the human's response + # Phase 2: Resume the retained in-memory run with the human's response. + # This response-only path requires the opaque token returned by Phase 1. + # Checkpoint restoration is a separate host-authorized path: checkpoint + # IDs locate persisted state but are not authorization credentials. print("\n=== Phase 2: Resume with feedback ===") print("(write_draft should NOT execute again — saved by @step)") result2 = await review_pipeline.run( From 7671ac79e38a60362b57a9c54f846c3c1828d8e1 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 12:48:17 +0900 Subject: [PATCH 4/7] Validate Python continuation hardening Run the complete Python workspace checks, aggregate coverage suite, repository hooks, and core package build from the final combined worktree. Keep the validation iteration code-neutral because all gates pass without corrective changes. Files changed: none; this commit records the final validation gate. Blockers: none. Next iteration: no remaining AFK tasks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> From 90ab95c6ee76213ad3da8ec5d6e7b9b24b53797d Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 13:49:04 +0900 Subject: [PATCH 5/7] Handle functional checkpoint continuation failures Publish retained continuation state only after checkpoint persistence succeeds, and cover reuse after a transient save failure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8f47743-1cdc-4924-8e1b-667d0d790b78 --- .../agent_framework/_workflows/_functional.py | 21 +++++++----- .../workflow/test_functional_workflow.py | 33 +++++++++++++++++++ 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 263f2bdeea1..088b71691c6 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -1101,11 +1101,20 @@ async def _on_step_completed() -> None: span.add_event(OtelAttr.WORKFLOW_COMPLETED) except WorkflowInterrupted: - # Persist step cache for response-only replay + pending_step_cache = dict(ctx._step_cache) + pending_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts) + pending_request_ids = set(ctx._pending_requests) + + # Persist before publishing in-memory continuation authority. If + # storage fails, the caller receives no token, so the workflow + # instance must remain free for a fresh run. + if storage is not None: + await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) + self._last_message = message - self._last_step_cache = dict(ctx._step_cache) - self._last_step_cache_auto_request_info_counts = dict(ctx._step_cache_auto_request_info_counts) - self._last_pending_request_ids = set(ctx._pending_requests) + self._last_step_cache = pending_step_cache + self._last_step_cache_auto_request_info_counts = pending_step_cache_auto_request_info_counts + self._last_pending_request_ids = pending_request_ids self._rotate_continuation_authority() result_continuation_token[0] = self._get_continuation_token() @@ -1118,10 +1127,6 @@ async def _on_step_completed() -> None: with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS) - # Save checkpoint - if storage is not None: - await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) - with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS) diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index c299c027002..78fe5331ddd 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -21,6 +21,7 @@ InMemoryCheckpointStorage, RunContext, StepWrapper, + WorkflowCheckpoint, WorkflowEvent, WorkflowRunResult, WorkflowRunState, @@ -716,6 +717,38 @@ async def wf(x: int, ctx: RunContext) -> str: class TestCheckpointing: + async def test_failed_pause_checkpoint_does_not_strand_in_memory_continuation(self): + class FailsFirstSaveStorage(InMemoryCheckpointStorage): + def __init__(self) -> None: + super().__init__() + self.save_attempts = 0 + + async def save(self, checkpoint: WorkflowCheckpoint) -> str: + self.save_attempts += 1 + if self.save_attempts == 1: + raise RuntimeError("checkpoint storage unavailable") + return await super().save(checkpoint) + + storage = FailsFirstSaveStorage() + + @workflow(checkpoint_storage=storage) + async def review(doc: str, ctx: RunContext) -> str: + feedback = await ctx.request_info(doc, response_type=str, request_id="review") + return f"{doc}:{feedback}" + + with pytest.raises(RuntimeError, match="checkpoint storage unavailable"): + await review.run("first") + + recovered = await review.run("second") + + assert recovered.continuation_token is not None + assert recovered.get_request_info_events()[0].data == "second" + completed = await review.run( + responses={"review": "approved"}, + continuation_token=recovered.continuation_token, + ) + assert completed.get_outputs() == ["second:approved"] + async def test_restored_checkpoint_issues_process_local_continuation_authority(self): storage = InMemoryCheckpointStorage() From 76d5f3b5e5cc5ad1f752e979f9977da081868d2a Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 14:43:43 +0900 Subject: [PATCH 6/7] Address functional continuation review findings Add owner recovery for lost tokens, harden malformed token validation, preserve consistent failure surfaces, and keep agent pending state aligned with resumable workflow state. Document process-local single-use continuation semantics and extend regression coverage across direct, streaming, checkpoint, and agent paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8f47743-1cdc-4924-8e1b-667d0d790b78 --- python/packages/core/AGENTS.md | 17 +- .../agent_framework/_workflows/_functional.py | 205 +++++++++++------- .../agent_framework/_workflows/_workflow.py | 5 +- .../workflow/test_functional_workflow.py | 181 ++++++++++++++++ 4 files changed, 324 insertions(+), 84 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 611cb499f0d..cfa9aee1f84 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -201,13 +201,16 @@ agent_framework/ - **`WorkflowRunResult`** - Non-streaming workflow result with Workflow Output `get_outputs()` and Intermediate Output `get_intermediate_outputs()` accessors. Functional workflows that pause for `request_info` also return an opaque, single-use `continuation_token`; pass it with `responses` for an in-memory - response-only resume. Each `FunctionalWorkflow` instance retains at most one in-memory continuation: callers must - resume it or call `abandon_continuation(token)` before starting new input or restoring a checkpoint. Use separate - workflow instances for independent in-memory runs; `FunctionalWorkflowAgent.abandon_continuation` delegates the - same operation. Checkpoint restoration is a separate host-authorized path: checkpoint IDs only locate persisted - state, while the host or storage adapter owns authorization and tenant isolation. A restored functional workflow - that pauses returns fresh process-local continuation authority. This does not alter graph-workflow request-info - authoritative resolution from PR #7500. + response-only resume. The token is process-local, is not a durable polling token, and is consumed before resumed + user code executes to prevent ambiguous failures from replaying side effects. Each `FunctionalWorkflow` instance + retains at most one in-memory continuation: callers must resume it or call `abandon_continuation(token)` before + starting new input or restoring a checkpoint. If the token is irretrievably lost, the workflow owner can use + `abandon_continuation(force=True)` to recover the instance; hosts must not expose forced abandonment to untrusted + callers. Use separate workflow instances for independent in-memory runs; `FunctionalWorkflowAgent` delegates the + same abandonment operations. Checkpoint restoration is a separate host-authorized path: checkpoint IDs only + locate persisted state, while the host or storage adapter owns authorization and tenant isolation. A restored + functional workflow that pauses returns fresh process-local continuation authority. This does not alter + graph-workflow request-info authoritative resolution from PR #7500. - **Orchestrators**: `SequentialOrchestrator`, `ConcurrentOrchestrator`, `GroupChatOrchestrator`, `MagenticOrchestrator`, `HandoffOrchestrator` ## Built-in Providers diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 088b71691c6..3dfc18af10a 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -660,6 +660,13 @@ class FunctionalWorkflow: edge wiring is involved. Native Python control flow (``if``/``else``, ``for``, ``asyncio.gather``) is used for branching and parallelism. + Continuation tokens for pending in-memory request-info work are + process-local, single-use capabilities. They must be supplied together + with ``responses`` and are consumed before resumed user code executes. + After an execution failure, recover from an authorized checkpoint or + start a new run after owner-authorized abandonment; reusing the consumed + token could otherwise duplicate side effects. + A workflow instance retains at most one in-memory continuation. Resume or explicitly abandon a pending continuation before starting new input or restoring a checkpoint on that instance. Use separate workflow @@ -824,6 +831,8 @@ def run( continuation_token: Opaque token returned by the immediately preceding result for the pending in-memory continuation. Required when *responses* are provided without *checkpoint_id*. + This process-local, single-use capability is consumed before + resumed user code executes and is not a durable polling token. checkpoint_id: Identifier of a checkpoint to restore from. Requires *checkpoint_storage* to be set (here or on the decorator). Checkpoint restoration does not use prior @@ -860,24 +869,16 @@ def run( "Cannot start or restore a functional workflow run while an in-memory continuation is pending. " "Resume or abandon the pending continuation first." ) - if responses is not None and checkpoint_id is None: - # Require at least one response key to match a currently-pending - # request; prevents silent replay against stale state while still - # allowing callers to accumulate prior answers across multi-round - # HITL. - if not self._last_pending_request_ids: - raise ValueError( - f"responses={list(responses)!r} do not correspond to any pending request on " - f"workflow '{self.name}'. The workflow has no pending request_info events, " - f"so there is nothing to resume. Start a fresh run with 'message', or supply " - f"'checkpoint_id' to restore a specific checkpoint." - ) - if not (set(responses) & self._last_pending_request_ids): - raise ValueError( - f"responses={list(responses)!r} do not answer any of the currently-pending " - f"requests on workflow '{self.name}' ({sorted(self._last_pending_request_ids)!r}). " - f"Provide a response keyed by one of the pending request_ids." - ) + # Require at least one response key to match a currently-pending + # request; prevents silent replay against stale state while still + # allowing callers to accumulate prior answers across multi-round + # HITL. + if responses is not None and checkpoint_id is None and not (set(responses) & self._last_pending_request_ids): + raise ValueError( + f"responses={list(responses)!r} do not answer any of the currently-pending " + f"requests on workflow '{self.name}' ({sorted(self._last_pending_request_ids)!r}). " + f"Provide a response keyed by one of the pending request_ids." + ) self._ensure_not_running() result_continuation_token: list[ContinuationToken | None] = [None] @@ -946,15 +947,26 @@ def as_agent( **kwargs, ) - def abandon_continuation(self, continuation_token: ContinuationToken | None = None) -> None: + def abandon_continuation( + self, + continuation_token: ContinuationToken | None = None, + *, + force: bool = False, + ) -> None: """Abandon the pending in-memory continuation. Successful abandonment consumes the token and clears the retained message, step cache, request metadata, and pending requests. A failed - attempt leaves the continuation unchanged. + token-authorized attempt leaves the continuation unchanged. + + ``force=True`` is an owner-only recovery escape hatch for a lost token. + Hosts must not expose forced abandonment to untrusted callers because + it allows one caller to cancel another caller's pending continuation. Args: continuation_token: Opaque token returned by the pending run. + force: Clear retained continuation state without validating a + token. Intended only for the owner of the workflow instance. Raises: RuntimeError: If the workflow is currently running. @@ -962,6 +974,9 @@ def abandon_continuation(self, continuation_token: ContinuationToken | None = No """ if self._is_running: raise RuntimeError("Cannot abandon a continuation while the functional workflow is running.") + if force: + self._clear_continuation_state() + return continuation_nonce = self._validate_continuation_authority(continuation_token) self._consume_continuation_authority(continuation_nonce) @@ -1109,7 +1124,12 @@ async def _on_step_completed() -> None: # storage fails, the caller receives no token, so the workflow # instance must remain free for a fresh run. if storage is not None: - await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) + try: + await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) + except Exception as exc: + for event in self._failure_events(ctx, span, exc): + yield event + raise self._last_message = message self._last_step_cache = pending_step_cache @@ -1133,26 +1153,29 @@ async def _on_step_completed() -> None: span.add_event(OtelAttr.WORKFLOW_COMPLETED) except Exception as exc: - # Yield any events collected before the failure - for event in ctx._get_events(): + for event in self._failure_events(ctx, span, exc): yield event - - details = WorkflowErrorDetails.from_exception(exc) - with _framework_event_origin(): - yield WorkflowEvent.failed(details) - with _framework_event_origin(): - yield WorkflowEvent.status(WorkflowRunState.FAILED) - - span.add_event( - name=OtelAttr.WORKFLOW_ERROR, - attributes={ - "error.message": str(exc), - "error.type": type(exc).__name__, - }, - ) - capture_exception(span, exception=exc) raise + @staticmethod + def _failure_events(ctx: RunContext, span: Any, exc: Exception) -> list[WorkflowEvent[Any]]: + events = ctx._get_events() + details = WorkflowErrorDetails.from_exception(exc) + with _framework_event_origin(): + events.append(WorkflowEvent.failed(details)) + with _framework_event_origin(): + events.append(WorkflowEvent.status(WorkflowRunState.FAILED)) + + span.add_event( + name=OtelAttr.WORKFLOW_ERROR, + attributes={ + "error.message": str(exc), + "error.type": type(exc).__name__, + }, + ) + capture_exception(span, exception=exc) + return events + async def _execute(self, ctx: RunContext, message: Any) -> Any: """Run the user's async function with the active context.""" if message is not None and not self._non_ctx_param_names: @@ -1309,22 +1332,32 @@ def _validate_continuation_authority(self, continuation_token: ContinuationToken if ( self._continuation_nonce is None or not isinstance(continuation_token, dict) - or set(continuation_token) != {"kind", "version", "token"} + or not {"kind", "version", "token"}.issubset(continuation_token) or continuation_token.get("kind") != _CONTINUATION_KIND or continuation_token.get("version") != _CONTINUATION_VERSION ): raise ValueError(_INVALID_CONTINUATION_AUTHORITY) token = continuation_token.get("token") - if not isinstance(token, str) or not secrets.compare_digest(token, self._continuation_nonce): + if not isinstance(token, str) or not self._continuation_tokens_equal(token, self._continuation_nonce): raise ValueError(_INVALID_CONTINUATION_AUTHORITY) return token def _consume_continuation_authority(self, continuation_nonce: str) -> None: - if self._continuation_nonce is None or not secrets.compare_digest(continuation_nonce, self._continuation_nonce): + if self._continuation_nonce is None or not self._continuation_tokens_equal( + continuation_nonce, + self._continuation_nonce, + ): raise ValueError(_INVALID_CONTINUATION_AUTHORITY) self._clear_continuation_state() + @staticmethod + def _continuation_tokens_equal(candidate: str, expected: str) -> bool: + try: + return secrets.compare_digest(candidate.encode(), expected.encode()) + except UnicodeEncodeError: + return False + def _clear_continuation_state(self) -> None: self._last_message = None self._last_step_cache = {} @@ -1452,8 +1485,10 @@ class FunctionalWorkflowAgent: response's ``continuation_token``; checkpoint restores use ``checkpoint_id=`` after the host or storage adapter authorizes access. A restored run that pauses returns fresh process-local authority through - the agent response's ``continuation_token``. :meth:`abandon_continuation` - delegates token-authorized abandonment to the wrapped workflow. + the agent response's ``continuation_token``. The token is not a durable + polling token and must be supplied together with ``responses``; providing + it alone does not resume work. :meth:`abandon_continuation` delegates + abandonment to the wrapped workflow. Args: workflow: The :class:`FunctionalWorkflow` to wrap. @@ -1491,9 +1526,14 @@ def pending_requests(self) -> dict[str, WorkflowEvent[Any]]: """Pending request_info events emitted during the last run.""" return self._pending_requests - def abandon_continuation(self, continuation_token: ContinuationToken | None = None) -> None: + def abandon_continuation( + self, + continuation_token: ContinuationToken | None = None, + *, + force: bool = False, + ) -> None: """Abandon the wrapped workflow's pending in-memory continuation.""" - self._workflow.abandon_continuation(continuation_token) + self._workflow.abandon_continuation(continuation_token, force=force) self._pending_requests = {} @overload @@ -1544,7 +1584,9 @@ def run( responses: HITL responses keyed by ``request_id``, forwarded to the underlying workflow so HITL resumes work via this agent. continuation_token: Opaque continuation token returned by the - preceding agent response. + preceding agent response. This process-local, single-use + capability is valid only together with *responses* and is not + a durable polling token. checkpoint_id: Optional host-authorized checkpoint to restore from. A checkpoint ID locates state; it is not an authorization credential. @@ -1584,7 +1626,7 @@ async def _run_non_streaming( checkpoint_storage: CheckpointStorage | None = None, **kwargs: Any, ) -> AgentResponse: - result = await self._workflow.run( + workflow_result = self._workflow.run( messages, responses=responses, continuation_token=continuation_token, @@ -1592,6 +1634,14 @@ async def _run_non_streaming( checkpoint_storage=checkpoint_storage, **kwargs, ) + # Synchronous validation has succeeded, so the prior pending-request + # view no longer describes the accepted run. + self._pending_requests = {} + try: + result = await workflow_result + except Exception: + self._pending_requests = {} + raise return self._result_to_agent_response(result) def _run_streaming( @@ -1607,8 +1657,6 @@ def _run_streaming( from .._types import Content agent_name = self.name - # Clear per-run pending state up front - self._pending_requests = {} workflow_stream = self._workflow.run( messages, stream=True, @@ -1618,34 +1666,41 @@ def _run_streaming( checkpoint_storage=checkpoint_storage, **kwargs, ) + # Synchronous workflow validation has succeeded, so this run now owns + # the adapter's pending-request view. + self._pending_requests = {} async def _generate_updates() -> AsyncIterable[AgentResponseUpdate]: - async for event in workflow_stream: - if event.type == "output": - data = event.data - if isinstance(data, str): - contents: list[Content] = [Content.from_text(text=data)] - elif isinstance(data, Content): - contents = [data] - else: - contents = [Content.from_text(text=str(data))] - yield AgentResponseUpdate( - contents=contents, - role="assistant", - author_name=agent_name, - ) - elif event.type == "request_info": - approval = self._request_info_to_approval_request(event) - if approval is None: - continue - yield AgentResponseUpdate( - contents=[approval], - role="assistant", - author_name=agent_name, - ) - workflow_result = await workflow_stream.get_final_response() - if workflow_result.continuation_token is not None: - yield AgentResponseUpdate(continuation_token=workflow_result.continuation_token) + try: + async for event in workflow_stream: + if event.type == "output": + data = event.data + if isinstance(data, str): + contents: list[Content] = [Content.from_text(text=data)] + elif isinstance(data, Content): + contents = [data] + else: + contents = [Content.from_text(text=str(data))] + yield AgentResponseUpdate( + contents=contents, + role="assistant", + author_name=agent_name, + ) + elif event.type == "request_info": + approval = self._request_info_to_approval_request(event) + if approval is None: + continue + yield AgentResponseUpdate( + contents=[approval], + role="assistant", + author_name=agent_name, + ) + workflow_result = await workflow_stream.get_final_response() + if workflow_result.continuation_token is not None: + yield AgentResponseUpdate(continuation_token=workflow_result.continuation_token) + except Exception: + self._pending_requests = {} + raise return ResponseStream( _generate_updates(), diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index dec9258028e..05f7784c67a 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -120,8 +120,9 @@ class WorkflowRunResult(list[WorkflowEvent]): - status_timeline(): Access the complete status event history Functional workflows set ``continuation_token`` when execution pauses for - external input. Callers must treat it as opaque and pass it back for the - next response-only in-memory resume. + external input. It is a process-local, single-use capability rather than a + durable polling token. Callers must treat it as opaque and pass it back + together with responses for the next in-memory resume. """ def __init__( diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 78fe5331ddd..24f3d7fbe13 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -283,6 +283,81 @@ async def private_wf(message: str, ctx: RunContext) -> str: ) assert completed.get_outputs() == ["caller-secret:authorized"] + async def test_non_ascii_continuation_token_is_rejected_as_invalid_authority(self): + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + return await ctx.request_info(doc, response_type=str, request_id="review") + + paused = await review_wf.run("draft") + malformed_token = json.loads(json.dumps(paused.continuation_token)) + malformed_token["token"] = "caf\u00e9" + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await review_wf.run( + responses={"review": "approved"}, + continuation_token=malformed_token, + ) + + async def test_unpaired_surrogate_continuation_token_is_rejected_as_invalid_authority(self): + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + return await ctx.request_info(doc, response_type=str, request_id="review") + + paused = await review_wf.run("draft") + malformed_token = json.loads(json.dumps(paused.continuation_token)) + malformed_token["token"] = json.loads(r'"\ud800"') + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await review_wf.run( + responses={"review": "approved"}, + continuation_token=malformed_token, + ) + + async def test_invalid_authority_does_not_reveal_whether_continuation_is_pending(self): + @workflow + async def idle_wf(value: str) -> str: + return value + + @workflow + async def pending_wf(value: str, ctx: RunContext) -> str: + return await ctx.request_info(value, response_type=str, request_id="review") + + await idle_wf.run("done") + paused = await pending_wf.run("draft") + invalid_token = json.loads(json.dumps(paused.continuation_token)) + invalid_token["token"] = "invalid" + + errors: list[str] = [] + for workflow_instance in (idle_wf, pending_wf): + with pytest.raises(ValueError) as exc_info: + await workflow_instance.run( + responses={"review": "approved"}, + continuation_token=invalid_token, + ) + errors.append(str(exc_info.value)) + + assert errors == [ + "Invalid functional workflow continuation authority.", + "Invalid functional workflow continuation authority.", + ] + + async def test_continuation_token_allows_additive_opaque_fields(self): + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + feedback = await ctx.request_info(doc, response_type=str, request_id="review") + return f"{doc}:{feedback}" + + paused = await review_wf.run("draft") + extended_token = json.loads(json.dumps(paused.continuation_token)) + extended_token["future_field"] = {"opaque": True} + + completed = await review_wf.run( + responses={"review": "approved"}, + continuation_token=extended_token, + ) + + assert completed.get_outputs() == ["draft:approved"] + async def test_request_info_interrupts(self): @workflow async def review_wf(doc: str, ctx: RunContext) -> str: @@ -408,6 +483,18 @@ async def review_wf(doc: str, ctx: RunContext) -> str: assert fresh_request.data == "prepared:new" assert step_calls == 2 + async def test_force_abandon_continuation_recovers_when_token_is_lost(self) -> None: + @workflow + async def review_wf(doc: str, ctx: RunContext) -> str: + return await ctx.request_info(doc, response_type=str, request_id="review") + + await review_wf.run("abandoned") + + review_wf.abandon_continuation(force=True) + fresh = await review_wf.run("fresh") + + assert fresh.get_request_info_events()[0].data == "fresh" + async def test_responses_while_pending_requests_does_not_warn(self, caplog: pytest.LogCaptureFixture) -> None: """Delivering responses is the normal completion path and must not warn.""" @@ -717,6 +804,27 @@ async def wf(x: int, ctx: RunContext) -> str: class TestCheckpointing: + async def test_failed_pause_checkpoint_uses_normal_failure_event_surface(self): + class FailingStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> str: + raise RuntimeError("checkpoint storage unavailable") + + storage = FailingStorage() + + @workflow(checkpoint_storage=storage) + async def review(doc: str, ctx: RunContext) -> str: + return await ctx.request_info(doc, response_type=str, request_id="review") + + events: list[WorkflowEvent] = [] + with pytest.raises(RuntimeError, match="checkpoint storage unavailable"): + async for event in review.run("draft", stream=True): + events.append(event) + + assert any(event.type == "request_info" for event in events) + assert any(event.type == "failed" for event in events) + assert events[-1].type == "status" + assert events[-1].state == WorkflowRunState.FAILED + async def test_failed_pause_checkpoint_does_not_strand_in_memory_continuation(self): class FailsFirstSaveStorage(InMemoryCheckpointStorage): def __init__(self) -> None: @@ -2094,6 +2202,58 @@ async def wf(x: str, ctx: RunContext) -> str: assert completed.text == "got:answered" assert completed.continuation_token is None + async def test_failed_streaming_resume_preserves_pending_requests(self): + @workflow + async def wf(x: str, ctx: RunContext) -> str: + return await ctx.request_info(x, response_type=str, request_id="rid-1") + + agent = wf.as_agent() + paused = await agent.run("topic") + wrong_token = json.loads(json.dumps(paused.continuation_token)) + wrong_token["token"] = "wrong" + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + agent.run( + responses={"rid-1": "answered"}, + continuation_token=wrong_token, + stream=True, + ) + + assert "rid-1" in agent.pending_requests + + async def test_failed_non_streaming_resume_clears_consumed_pending_request(self): + @workflow + async def wf(x: str, ctx: RunContext) -> str: + answer = await ctx.request_info(x, response_type=str, request_id="rid-1") + raise RuntimeError(f"resume failed after {answer}") + + agent = wf.as_agent() + paused = await agent.run("topic") + + with pytest.raises(RuntimeError, match="resume failed after answered"): + await agent.run( + responses={"rid-1": "answered"}, + continuation_token=paused.continuation_token, + ) + + assert agent.pending_requests == {} + + async def test_failed_pause_checkpoint_does_not_leave_agent_pending_request(self): + class FailingStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> str: + raise RuntimeError("checkpoint storage unavailable") + + @workflow(checkpoint_storage=FailingStorage()) + async def wf(x: str, ctx: RunContext) -> str: + return await ctx.request_info(x, response_type=str, request_id="rid-1") + + agent = wf.as_agent() + + with pytest.raises(RuntimeError, match="checkpoint storage unavailable"): + await agent.run("topic", stream=True).get_final_response() + + assert agent.pending_requests == {} + async def test_agent_can_abandon_pending_continuation(self) -> None: @workflow async def wf(x: str, ctx: RunContext) -> str: @@ -2122,6 +2282,20 @@ async def wf(x: str, ctx: RunContext) -> str: assert fresh.continuation_token is not None assert fresh.continuation_token != abandoned.continuation_token + async def test_agent_can_force_abandon_when_continuation_token_is_lost(self) -> None: + @workflow + async def wf(x: str, ctx: RunContext) -> str: + return await ctx.request_info(x, response_type=str, request_id="rid-1") + + agent = wf.as_agent() + await agent.run("abandoned") + + agent.abandon_continuation(force=True) + + assert agent.pending_requests == {} + fresh = await agent.run("fresh") + assert fresh.continuation_token is not None + class TestRunDocstringAllowsResponsesAndCheckpoint: """Regression for bug_010: docstring must permit responses+checkpoint_id combo.""" @@ -2131,6 +2305,13 @@ def test_docstring_says_at_least_one(self): assert "At least one" in doc or "at least one" in doc assert "Exactly one" not in doc + def test_agent_docstring_distinguishes_process_local_token_from_durable_polling(self): + doc = " ".join((FunctionalWorkflowAgent.__doc__ or "").split()) + + assert "process-local" in doc + assert "not a durable polling token" in doc + assert "must be supplied together with" in doc + class TestFunctionalWorkflowExperimentalStage: """Tests for the experimental stage annotations applied to functional workflow APIs.""" From b128713b4ad497e5b180cd633ad4dff5f0f3caa5 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 15:09:57 +0900 Subject: [PATCH 7/7] Handle functional continuation cancellation Release the workflow run guard when cancellation interrupts resumed user code while keeping the single-use continuation token consumed. Replace sample assertions with explicit runtime checks and add cancellation regression coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8f47743-1cdc-4924-8e1b-667d0d790b78 --- .../agent_framework/_workflows/_functional.py | 5 +++ .../workflow/test_functional_workflow.py | 35 +++++++++++++++++++ .../03-workflows/functional/hitl_review.py | 9 +++-- 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 3dfc18af10a..ccfff2b0a25 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -36,6 +36,7 @@ # pyright: reportPrivateUsage=false # Classes in this module (RunContext, StepWrapper, FunctionalWorkflow) form a # cohesive unit and intentionally access each other's underscore-prefixed members. +import asyncio import functools import hashlib import inspect @@ -1152,6 +1153,10 @@ async def _on_step_completed() -> None: span.add_event(OtelAttr.WORKFLOW_COMPLETED) + except asyncio.CancelledError: + await self._run_cleanup() + raise + except Exception as exc: for event in self._failure_events(ctx, span, exc): yield event diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 24f3d7fbe13..e04e6fe7db3 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -600,6 +600,41 @@ async def failing_resume(data: str, ctx: RunContext) -> str: assert fresh.continuation_token is not None assert fresh.get_request_info_events()[0].data == "fresh" + async def test_cancellation_after_token_consumption_releases_workflow_instance(self): + resumed_user_code_started = asyncio.Event() + keep_running = asyncio.Event() + + @workflow + async def cancellable_resume(data: str, ctx: RunContext) -> str: + answer = await ctx.request_info(data, response_type=str, request_id="r1") + resumed_user_code_started.set() + await keep_running.wait() + return f"{data}:{answer}" + + paused = await cancellable_resume.run("input") + + async def resume() -> None: + await cancellable_resume.run( + responses={"r1": "response"}, + continuation_token=paused.continuation_token, + ) + + task = asyncio.create_task(resume()) + await resumed_user_code_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + await cancellable_resume.run( + responses={"r1": "response"}, + continuation_token=paused.continuation_token, + ) + + fresh = await cancellable_resume.run("fresh") + assert fresh.continuation_token is not None + async def test_request_info_auto_generates_id(self): @workflow async def auto_id_wf(x: int, ctx: RunContext) -> None: diff --git a/python/samples/03-workflows/functional/hitl_review.py b/python/samples/03-workflows/functional/hitl_review.py index bd360854924..08c7ea2818f 100644 --- a/python/samples/03-workflows/functional/hitl_review.py +++ b/python/samples/03-workflows/functional/hitl_review.py @@ -66,11 +66,14 @@ async def main(): # If request_info() was reached, the state is IDLE_WITH_PENDING_REQUESTS. # If the workflow completed without hitting request_info(), it would be IDLE. print(f"State: {(final_state := result1.get_final_state())}") - assert final_state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + if final_state != WorkflowRunState.IDLE_WITH_PENDING_REQUESTS: + raise RuntimeError(f"Expected pending review input, but workflow entered {final_state}.") requests = result1.get_request_info_events() print(f"Pending request: {requests[0].request_id}") - assert result1.continuation_token is not None + continuation_token = result1.continuation_token + if continuation_token is None: + raise RuntimeError("Expected a continuation token for the pending review.") # Phase 2: Resume the retained in-memory run with the human's response. # This response-only path requires the opaque token returned by Phase 1. @@ -80,7 +83,7 @@ async def main(): print("(write_draft should NOT execute again — saved by @step)") result2 = await review_pipeline.run( responses={"review_request": "Add more details about alignment research"}, - continuation_token=result1.continuation_token, + continuation_token=continuation_token, ) print(f"State: {result2.get_final_state()}")