From fbb01312de87a32a7e8632615e5cf97d653bf8c4 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 5 Aug 2026 12:18:31 +0900 Subject: [PATCH 01/10] 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 02/10] 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 03/10] 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 04/10] 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 05/10] 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 06/10] 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 07/10] 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()}") From 9082f1788ad0d1fd30b49ade13124274f8766a52 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 12 Aug 2026 14:20:14 +0900 Subject: [PATCH 08/10] Simplify functional workflow instance isolation Remove continuation-token handling and align functional workflows with the graph workflow ownership model: one stateful instance per logical caller or session. Add create_instance() for independent callers, document the ownership contract, and cover pending-state isolation between instances. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8f47743-1cdc-4924-8e1b-667d0d790b78 --- python/packages/core/AGENTS.md | 18 +- .../agent_framework/_workflows/_checkpoint.py | 8 +- .../agent_framework/_workflows/_functional.py | 431 ++++------- .../agent_framework/_workflows/_workflow.py | 15 +- .../workflow/test_functional_workflow.py | 691 ++---------------- .../03-workflows/functional/hitl_review.py | 17 +- 6 files changed, 196 insertions(+), 984 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index cfa9aee1f84..478b6e941b5 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -199,18 +199,12 @@ 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. 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. 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. + and Intermediate Output `get_intermediate_outputs()` accessors +- **`FunctionalWorkflow` instance ownership** - Functional workflows retain the original message, cached step + results, and pending `request_info` state on the workflow instance for response-only replay. Like graph workflows, + one instance represents one logical caller or session. Do not share an instance or its + `FunctionalWorkflowAgent` across mutually untrusted callers. Use `create_instance()` to create independent + stateful instances from the same decorated workflow definition. - **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 3dc600cb915..2b267979e99 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -127,13 +127,7 @@ def from_dict(cls, data: Mapping[str, Any]) -> WorkflowCheckpoint: class CheckpointStorage(Protocol): - """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. - """ + """Protocol for checkpoint storage backends.""" 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 ccfff2b0a25..880b4de7ae1 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -36,12 +36,10 @@ # 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 import logging -import secrets import typing from collections.abc import AsyncIterable, Awaitable, Callable, Sequence from contextvars import ContextVar @@ -50,7 +48,7 @@ from .._feature_stage import ExperimentalFeature, experimental from .._serialization import make_json_safe -from .._types import AgentResponse, AgentResponseUpdate, ContinuationToken, ResponseStream +from .._types import AgentResponse, AgentResponseUpdate, ResponseStream from ..observability import OtelAttr, capture_exception, create_workflow_span from ._checkpoint import CheckpointStorage, WorkflowCheckpoint from ._events import ( @@ -65,17 +63,6 @@ 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) @@ -218,10 +205,8 @@ 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}, - continuation_token=prior_result.continuation_token)``, the same - function re-executes and ``request_info`` returns the provided *value* - directly. + ``run(responses={request_id: value})``, the same function re-executes + and ``request_info`` returns the provided *value* directly. Args: request_data: Arbitrary payload describing what information is @@ -661,28 +646,11 @@ 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 - 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. + Like graph-based :class:`Workflow`, each instance owns mutable execution + state across calls to :meth:`run`. Scope an instance to one logical + caller or session. Use :meth:`create_instance` when independent callers + need the same workflow definition; do not share one instance across + mutually untrusted callers. Args: func: The async function that implements the workflow logic. @@ -725,7 +693,6 @@ 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) @@ -738,6 +705,22 @@ def __init__( functools.update_wrapper(self, func) # type: ignore[arg-type] + def create_instance(self) -> FunctionalWorkflow: + """Create an independent stateful instance of this workflow definition. + + The new instance reuses the workflow function and configuration but + has its own message, step-cache, and pending-request state. + + Returns: + A new :class:`FunctionalWorkflow` for one logical caller or session. + """ + return FunctionalWorkflow( + self._func, + name=self.name, + description=self.description, + checkpoint_storage=self._checkpoint_storage, + ) + @staticmethod def _classify_signature(func: Callable[..., Any]) -> list[str]: """Return the names of non-ctx parameters, validating arity. @@ -779,7 +762,6 @@ 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, @@ -792,7 +774,6 @@ 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, @@ -805,7 +786,6 @@ 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, @@ -829,17 +809,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 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 - process-local continuation authority; the host or storage - adapter is responsible for authorizing checkpoint access and - tenant isolation. + decorator). checkpoint_storage: Override the default checkpoint storage for this run. include_status_events: When ``True`` (non-streaming only), @@ -857,48 +829,58 @@ def run( Raises: ValueError: If the combination of *message*, *responses*, and *checkpoint_id* is invalid. - RuntimeError: If the workflow is already running, or if new input - or a checkpoint restore is attempted while an in-memory - continuation is pending. + RuntimeError: If the workflow is already running (concurrent + 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. if (message is not None or checkpoint_id is not None) and self._last_pending_request_ids: - raise RuntimeError( - "Cannot start or restore a functional workflow run while an in-memory continuation is pending. " - "Resume or abandon the pending continuation first." - ) - # 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." + 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" + ), ) + if responses 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." + ) 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, - continuation_token=result_continuation_token, - ), + finalizer=functools.partial(self._finalize_events, include_status_events=include_status_events), cleanup_hooks=[self._run_cleanup], ) @@ -948,39 +930,6 @@ def as_agent( **kwargs, ) - 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 - 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. - 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.") - if force: - self._clear_continuation_state() - return - continuation_nonce = self._validate_continuation_authority(continuation_token) - self._consume_continuation_authority(continuation_nonce) - # ------------------------------------------------------------------ # Internal execution # ------------------------------------------------------------------ @@ -990,8 +939,6 @@ 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, @@ -1038,6 +985,10 @@ 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) @@ -1048,7 +999,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, message, ckpt_chain[0]) + ckpt_chain[0] = await self._save_checkpoint(ctx, storage, ckpt_chain[0]) ctx._on_step_completed = _on_step_completed @@ -1067,9 +1018,6 @@ 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) @@ -1098,46 +1046,29 @@ async def _on_step_completed() -> None: # Save final checkpoint if storage is available if storage is not None: - await self._save_checkpoint(ctx, storage, message, ckpt_chain[0]) + await self._save_checkpoint(ctx, storage, 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() with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS) else: # Clean completion — drop cross-run replay state. - self._clear_continuation_state() + self._last_message = None + self._last_step_cache = {} + self._last_step_cache_auto_request_info_counts = {} + self._last_pending_request_ids = set() with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE) span.add_event(OtelAttr.WORKFLOW_COMPLETED) except WorkflowInterrupted: - 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: - 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 - 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() + # Persist step cache for response-only replay + 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) # HITL interruption — yield events collected so far for event in ctx._get_events(): @@ -1148,38 +1079,35 @@ 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, ckpt_chain[0]) + with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS) 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 any events collected before the failure + for event in ctx._get_events(): yield event - 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 + 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 async def _execute(self, ctx: RunContext, message: Any) -> Any: """Run the user's async function with the active context.""" @@ -1237,13 +1165,12 @@ 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"] = original_message + state["_original_message"] = self._last_message checkpoint = WorkflowCheckpoint( workflow_name=self.name, @@ -1304,7 +1231,6 @@ 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]] = [] @@ -1319,56 +1245,7 @@ def _finalize_events( continue filtered.append(ev) - 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 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 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 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 = {} - self._last_step_cache_auto_request_info_counts = {} - self._last_pending_request_ids = set() - self._continuation_nonce = None + return WorkflowRunResult(filtered, status_events) @staticmethod def _validate_run_params( @@ -1486,14 +1363,11 @@ 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. Response-only callers resume via ``responses=`` and the prior - 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``. 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. + adapter. Callers resume via ``responses=`` / ``checkpoint_id=``. + + The wrapped workflow owns mutable execution state. Scope the workflow and + this adapter to one logical caller or session; create separate workflow + instances for independent or mutually untrusted callers. Args: workflow: The :class:`FunctionalWorkflow` to wrap. @@ -1531,16 +1405,6 @@ 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, - *, - force: bool = False, - ) -> None: - """Abandon the wrapped workflow's pending in-memory continuation.""" - self._workflow.abandon_continuation(continuation_token, force=force) - self._pending_requests = {} - @overload def run( self, @@ -1548,7 +1412,6 @@ 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, @@ -1561,7 +1424,6 @@ 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, @@ -1573,7 +1435,6 @@ 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, @@ -1588,13 +1449,7 @@ 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. 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. + checkpoint_id: Optional checkpoint to restore from. checkpoint_storage: Override the workflow's default :class:`CheckpointStorage` for this run. **kwargs: Extra keyword arguments forwarded to the workflow run. @@ -1607,7 +1462,6 @@ def run( return self._run_streaming( messages, responses=responses, - continuation_token=continuation_token, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, **kwargs, @@ -1615,7 +1469,6 @@ def run( return self._run_non_streaming( messages, responses=responses, - continuation_token=continuation_token, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, **kwargs, @@ -1626,27 +1479,17 @@ 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, ) -> AgentResponse: - workflow_result = self._workflow.run( + result = await self._workflow.run( messages, responses=responses, - continuation_token=continuation_token, checkpoint_id=checkpoint_id, 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( @@ -1654,7 +1497,6 @@ 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, @@ -1662,50 +1504,41 @@ 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, responses=responses, - continuation_token=continuation_token, checkpoint_id=checkpoint_id, 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]: - 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 + 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, + ) return ResponseStream( _generate_updates(), @@ -1761,4 +1594,4 @@ def _result_to_agent_response(self, result: WorkflowRunResult) -> AgentResponse: if approval_contents: messages.append(Msg("assistant", approval_contents)) - return AgentResponse(messages=messages, continuation_token=result.continuation_token) + return AgentResponse(messages=messages) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 05f7784c67a..77060b93a91 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 ContinuationToken, ResponseStream +from .._types import ResponseStream from ..exceptions import WorkflowException from ..observability import OtelAttr, capture_exception, create_workflow_span from ._checkpoint import CheckpointStorage @@ -118,22 +118,11 @@ 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. 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__( - self, - events: list[WorkflowEvent[Any]], - status_events: list[WorkflowEvent[Any]] | None = None, - continuation_token: ContinuationToken | None = None, - ) -> None: + def __init__(self, events: list[WorkflowEvent[Any]], status_events: list[WorkflowEvent[Any]] | 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 e04e6fe7db3..259699e4d6f 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -21,7 +21,6 @@ InMemoryCheckpointStorage, RunContext, StepWrapper, - WorkflowCheckpoint, WorkflowEvent, WorkflowRunResult, WorkflowRunState, @@ -229,134 +228,23 @@ async def par_wf(x: int) -> tuple[int, int]: class TestHITL: - async def test_response_only_resume_requires_returned_continuation_token(self): + async def test_separate_instances_isolate_pending_continuations(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_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") + feedback = await ctx.request_info(doc, response_type=str) 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} + caller_a = review_wf.create_instance() + caller_b = review_wf.create_instance() - completed = await review_wf.run( - responses={"review": "approved"}, - continuation_token=extended_token, - ) + caller_a_paused = await caller_a.run("caller-a") + request_id = caller_a_paused.get_request_info_events()[0].request_id - assert completed.get_outputs() == ["draft:approved"] + with pytest.raises(ValueError, match="no pending request_info events"): + await caller_b.run(responses={request_id: "caller-b-response"}) + + caller_a_completed = await caller_a.run(responses={request_id: "caller-a-response"}) + assert caller_a_completed.get_outputs() == ["caller-a:caller-a-response"] async def test_request_info_interrupts(self): @workflow @@ -382,118 +270,29 @@ 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!"}, - continuation_token=result1.continuation_token, - ) + result2 = await review_wf.run(responses={"req1": "Looks great!"}) outputs = result2.get_outputs() assert outputs == ["Final: Looks great!"] assert result2.get_final_state() == WorkflowRunState.IDLE - async def test_fresh_message_while_pending_requests_is_rejected_without_losing_continuation(self) -> None: + 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.""" + @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}" - 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") - - 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_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") + result1 = await review_wf.run("my doc") + assert result1.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS - review_wf.abandon_continuation(force=True) - fresh = await review_wf.run("fresh") + # 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): + await review_wf.run("another doc") - assert fresh.get_request_info_events()[0].data == "fresh" + assert "request_info event(s) are still pending" in caplog.text + assert "a fresh message" in caplog.text 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.""" @@ -508,10 +307,7 @@ 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!"}, - continuation_token=result1.continuation_token, - ) + result2 = await review_wf.run(responses={"req1": "Looks great!"}) assert result2.get_final_state() == WorkflowRunState.IDLE assert "still pending" not in caplog.text @@ -527,10 +323,7 @@ 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"}, - continuation_token=result1.continuation_token, - ) + result2 = await review_wf.run(responses={"req1": "LGTM"}) assert result2.get_outputs() == ["Final: LGTM"] async def test_multiple_sequential_interrupts(self): @@ -538,102 +331,21 @@ 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"{data}:{r1}+{r2}" + return f"{r1}+{r2}" # Phase 1: first interrupt 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"}, - continuation_token=result1.continuation_token, - ) + result2 = await multi_hitl.run(responses={"r1": "A"}) 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"}, - continuation_token=result2.continuation_token, - ) - 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): - 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, - ) - - 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_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 + result3 = await multi_hitl.run(responses={"r1": "A", "r2": "B"}) + assert result3.get_outputs() == ["A+B"] async def test_request_info_auto_generates_id(self): @workflow @@ -766,25 +478,6 @@ 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 @@ -839,121 +532,6 @@ 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: - 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() - - 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() @@ -1028,7 +606,6 @@ 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") @@ -1059,7 +636,6 @@ 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") @@ -1622,10 +1198,7 @@ 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"}, - continuation_token=result1.continuation_token, - ) + result2 = await wf.run(responses={"s1": "LGTM"}) assert result2.get_outputs() == ["reviewed: LGTM"] async def test_step_works_outside_workflow_with_explicit_ctx(self): @@ -1700,14 +1273,11 @@ async def wf(doc: str, ctx: RunContext) -> str: return f"got: {val}" # Phase 1 - paused = await wf.run("start") + 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}, - continuation_token=paused.continuation_token, - ) + result = await wf.run(responses={"r1": None}) assert result.get_outputs() == ["got: None"] assert any("None" in msg and "r1" in msg for msg in logs) @@ -1720,11 +1290,8 @@ 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}" - paused = await wf.run(1) - result = await wf.run( - responses={"r1": None}, - continuation_token=paused.continuation_token, - ) + await wf.run(1) + result = await wf.run(responses={"r1": None}) assert result.get_outputs() == ["value=None"] @@ -1784,10 +1351,7 @@ 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"}, - continuation_token=result1.continuation_token, - ) + result2 = await wf.run(responses={"r1": "ok"}) assert call_count_a == 1 # step_a not called again assert result2.get_outputs() == ["6:ok"] @@ -1817,10 +1381,7 @@ 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"}, - continuation_token=result1.continuation_token, - ) + result2 = await wf.run(responses={"rev": "LGTM"}) assert result2.get_outputs() == ["reviewed(30):LGTM"] # Phase 3: restore from latest checkpoint -- both steps should be bypassed @@ -1845,13 +1406,10 @@ async def needs_feedback(doc: str, ctx: RunContext) -> str: async def wf(doc: str) -> str: return await needs_feedback(doc) - paused = await wf.run("draft") + await wf.run("draft") with caplog_context(logging.getLogger("agent_framework._workflows._functional")) as logs: - result = await wf.run( - responses={"r1": None}, - continuation_token=paused.continuation_token, - ) + result = await wf.run(responses={"r1": None}) assert result.get_outputs() == ["got:None"] assert any("None" in msg and "r1" in msg for msg in logs) @@ -1896,10 +1454,7 @@ 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"}, - continuation_token=result1.continuation_token, - ) + result2 = await wf.run(responses={rid: "hello"}) assert result2.get_final_state() == WorkflowRunState.IDLE assert result2.get_outputs() == ["got:hello"] @@ -1912,16 +1467,10 @@ 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"}, - continuation_token=r1.continuation_token, - ) + r2 = await wf.run(responses={rid1: "A"}) rid2 = r2.get_request_info_events()[0].request_id assert rid1 != rid2 - r3 = await wf.run( - responses={rid1: "A", rid2: "B"}, - continuation_token=r2.continuation_token, - ) + r3 = await wf.run(responses={rid1: "A", rid2: "B"}) assert r3.get_outputs() == ["A/B"] async def test_cached_step_advances_auto_request_id_counter(self): @@ -1947,18 +1496,12 @@ 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"}, - continuation_token=first_run.continuation_token, - ) + second_run = await wf.run(responses={first_request_id: "A"}) 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"}, - continuation_token=second_run.continuation_token, - ) + final_run = await wf.run(responses={first_request_id: "A", second_request_id: "B"}) assert call_count == completed_call_count assert final_run.get_outputs() == ["A/B"] @@ -1976,15 +1519,9 @@ async def wf(x: int, ctx: RunContext) -> str: b = await ctx.request_info("q2", response_type=str, request_id="r2") return f"{a}/{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, - ) + await wf.run(1) + await wf.run(responses={"r1": "A"}) + result = await wf.run(responses={"r1": "A", "r2": "B"}) assert result.get_final_state() == WorkflowRunState.IDLE # Latest checkpoint must show no pending requests. checkpoints = await storage.list_checkpoints(workflow_name="wf") @@ -2032,7 +1569,7 @@ async def wf(x: int) -> int: return x * 2 await wf.run(5) # clean completion, no pending requests - with pytest.raises(ValueError, match="Invalid functional workflow continuation authority"): + with pytest.raises(ValueError, match="no pending request_info"): await wf.run(responses={"stale": "x"}) async def test_responses_mismatched_key_raises(self): @@ -2040,12 +1577,9 @@ 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") - paused = await wf.run(1) # interrupts with r1 pending + await wf.run(1) # interrupts with r1 pending with pytest.raises(ValueError, match="do not answer"): - await wf.run( - responses={"definitely_not_r1": "x"}, - continuation_token=paused.continuation_token, - ) + await wf.run(responses={"definitely_not_r1": "x"}) class TestReservedStateKeys: @@ -2203,13 +1737,9 @@ async def wf(x: str, ctx: RunContext) -> str: agent = wf.as_agent() # First phase: suspend - paused = await agent.run("topic") - assert paused.continuation_token is not None + await agent.run("topic") # Second phase: resume via the agent surface - response = await agent.run( - responses={"rid-1": "answered"}, - continuation_token=paused.continuation_token, - ) + response = await agent.run(responses={"rid-1": "answered"}) # Agent's final response should contain the workflow's text output. text_blobs: list[str] = [] for message in response.messages: @@ -2219,118 +1749,6 @@ 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 - - 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: - 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 - - 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.""" @@ -2340,13 +1758,6 @@ 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.""" diff --git a/python/samples/03-workflows/functional/hitl_review.py b/python/samples/03-workflows/functional/hitl_review.py index 08c7ea2818f..acca1974a8f 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 the returned continuation token. +external input and resuming with run(responses={...}). 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). @@ -71,20 +71,11 @@ async def main(): requests = result1.get_request_info_events() print(f"Pending request: {requests[0].request_id}") - 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. - # Checkpoint restoration is a separate host-authorized path: checkpoint - # IDs locate persisted state but are not authorization credentials. + + # 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"}, - continuation_token=continuation_token, - ) + result2 = await review_pipeline.run(responses={"review_request": "Add more details about alignment research"}) print(f"State: {result2.get_final_state()}") print(f"Output: {result2.get_outputs()[0]}") From 98f2bb8423e9219c464a63eb61141ce37173525c Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 12 Aug 2026 14:56:22 +0900 Subject: [PATCH 09/10] Scope functional workflow checkpoint storage Do not inherit checkpoint storage when creating an independent workflow instance. Allow hosts to provide an explicitly caller-scoped storage adapter and document that shared checkpoint access requires host authorization and tenant isolation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8f47743-1cdc-4924-8e1b-667d0d790b78 --- python/packages/core/AGENTS.md | 4 ++- .../agent_framework/_workflows/_functional.py | 24 +++++++++++++---- .../workflow/test_functional_workflow.py | 26 +++++++++++++++++++ 3 files changed, 48 insertions(+), 6 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 478b6e941b5..e3e2d49a11b 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -204,7 +204,9 @@ agent_framework/ results, and pending `request_info` state on the workflow instance for response-only replay. Like graph workflows, one instance represents one logical caller or session. Do not share an instance or its `FunctionalWorkflowAgent` across mutually untrusted callers. Use `create_instance()` to create independent - stateful instances from the same decorated workflow definition. + stateful instances from the same decorated workflow definition. New instances do not inherit the definition's + checkpoint storage; pass a caller-scoped storage explicitly when needed. Instance separation protects in-memory + state only — hosts remain responsible for authorizing and tenant-scoping access to any shared checkpoint adapter. - **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 880b4de7ae1..7c38ef1acf4 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -650,7 +650,9 @@ class FunctionalWorkflow: state across calls to :meth:`run`. Scope an instance to one logical caller or session. Use :meth:`create_instance` when independent callers need the same workflow definition; do not share one instance across - mutually untrusted callers. + mutually untrusted callers. Instance isolation covers in-memory state. + When checkpointing is enabled, the host must also authorize and + tenant-scope access to the checkpoint storage. Args: func: The async function that implements the workflow logic. @@ -705,11 +707,21 @@ def __init__( functools.update_wrapper(self, func) # type: ignore[arg-type] - def create_instance(self) -> FunctionalWorkflow: + def create_instance( + self, + *, + checkpoint_storage: CheckpointStorage | None = None, + ) -> FunctionalWorkflow: """Create an independent stateful instance of this workflow definition. The new instance reuses the workflow function and configuration but - has its own message, step-cache, and pending-request state. + has its own message, step-cache, and pending-request state. Checkpoint + storage is not inherited because a shared external store requires + host-managed authorization and tenant isolation. + + Args: + checkpoint_storage: Optional storage scoped and authorized for the + new instance's logical caller or session. Returns: A new :class:`FunctionalWorkflow` for one logical caller or session. @@ -718,7 +730,7 @@ def create_instance(self) -> FunctionalWorkflow: self._func, name=self.name, description=self.description, - checkpoint_storage=self._checkpoint_storage, + checkpoint_storage=checkpoint_storage, ) @staticmethod @@ -1367,7 +1379,9 @@ class FunctionalWorkflowAgent: The wrapped workflow owns mutable execution state. Scope the workflow and this adapter to one logical caller or session; create separate workflow - instances for independent or mutually untrusted callers. + instances for independent or mutually untrusted callers. If those + instances use checkpoint storage, the host must also authorize and + tenant-scope access to that external store. Args: workflow: The :class:`FunctionalWorkflow` to wrap. diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 259699e4d6f..1ddfc1954a7 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -246,6 +246,32 @@ async def review_wf(doc: str, ctx: RunContext) -> str: caller_a_completed = await caller_a.run(responses={request_id: "caller-a-response"}) assert caller_a_completed.get_outputs() == ["caller-a:caller-a-response"] + async def test_create_instance_does_not_inherit_checkpoint_storage_by_default(self): + storage = InMemoryCheckpointStorage() + + @workflow(checkpoint_storage=storage) + async def review_wf(doc: str, ctx: RunContext) -> str: + return await ctx.request_info(doc, response_type=str) + + caller = review_wf.create_instance() + await caller.run("caller") + + assert await storage.list_checkpoints(workflow_name="review_wf") == [] + + async def test_create_instance_accepts_tenant_scoped_checkpoint_storage(self): + definition_storage = InMemoryCheckpointStorage() + caller_storage = InMemoryCheckpointStorage() + + @workflow(checkpoint_storage=definition_storage) + async def review_wf(doc: str, ctx: RunContext) -> str: + return await ctx.request_info(doc, response_type=str) + + caller = review_wf.create_instance(checkpoint_storage=caller_storage) + await caller.run("caller") + + assert await definition_storage.list_checkpoints(workflow_name="review_wf") == [] + assert len(await caller_storage.list_checkpoints(workflow_name="review_wf")) == 1 + async def test_request_info_interrupts(self): @workflow async def review_wf(doc: str, ctx: RunContext) -> str: From 5a7508ea2bed3face4af1683530e3f3e1772a9db Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 12 Aug 2026 15:41:59 +0900 Subject: [PATCH 10/10] Require building functional workflow instances Make @workflow return a stateless FunctionalWorkflowDefinition and require build() before run() or as_agent(). This aligns functional workflows with the graph definition/build lifecycle and prevents module-level decorated definitions from retaining caller state. Move checkpoint configuration to build(), export the definition type, migrate samples, and cover isolated built instances. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8f47743-1cdc-4924-8e1b-667d0d790b78 --- python/PACKAGE_STATUS.md | 3 +- python/packages/core/AGENTS.md | 13 +- .../packages/core/agent_framework/__init__.py | 2 + .../core/agent_framework/__init__.pyi | 2 + .../agent_framework/_workflows/_functional.py | 117 +++++---- .../workflow/test_functional_workflow.py | 246 ++++++++++-------- .../05_functional_workflow_with_agents.py | 3 +- .../06_functional_workflow_basics.py | 3 +- .../functional/agent_integration.py | 7 +- .../03-workflows/functional/basic_pipeline.py | 8 +- .../functional/basic_streaming_pipeline.py | 10 +- .../03-workflows/functional/hitl_review.py | 6 +- .../functional/naive_group_chat.py | 3 +- .../functional/parallel_pipeline.py | 9 +- .../functional/steps_and_checkpointing.py | 13 +- 15 files changed, 253 insertions(+), 192 deletions(-) diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index 37053838277..0c7e987ff06 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -102,7 +102,8 @@ listed below. - `agent-framework-core`: functional workflow APIs from `agent_framework/_workflows/_functional.py`, including `RunContext`, `step`, - `FunctionalWorkflow`, `workflow`, and `FunctionalWorkflowAgent` + `FunctionalWorkflowDefinition`, `FunctionalWorkflow`, `workflow`, and + `FunctionalWorkflowAgent` #### `HARNESS` diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index e3e2d49a11b..2b65468b9fd 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -200,13 +200,12 @@ agent_framework/ 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 -- **`FunctionalWorkflow` instance ownership** - Functional workflows retain the original message, cached step - results, and pending `request_info` state on the workflow instance for response-only replay. Like graph workflows, - one instance represents one logical caller or session. Do not share an instance or its - `FunctionalWorkflowAgent` across mutually untrusted callers. Use `create_instance()` to create independent - stateful instances from the same decorated workflow definition. New instances do not inherit the definition's - checkpoint storage; pass a caller-scoped storage explicitly when needed. Instance separation protects in-memory - state only — hosts remain responsible for authorizing and tenant-scoping access to any shared checkpoint adapter. +- **Functional workflow definition/build lifecycle** - `@workflow` returns a stateless + `FunctionalWorkflowDefinition`. Call `build()` to create a stateful `FunctionalWorkflow` scoped to one logical + caller or session. The definition has no `run()` or `as_agent()` surface, so module-level decorated definitions + cannot accidentally retain caller state. Each built workflow and its `FunctionalWorkflowAgent` must remain scoped + to that caller/session. Pass a caller-scoped checkpoint storage to `build(checkpoint_storage=...)` when needed; + hosts remain responsible for authorizing and tenant-scoping access to any shared checkpoint adapter. - **Orchestrators**: `SequentialOrchestrator`, `ConcurrentOrchestrator`, `GroupChatOrchestrator`, `MagenticOrchestrator`, `HandoffOrchestrator` ## Built-in Providers diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ff0aa20e1a0..528ad5e1740 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -319,6 +319,7 @@ "._workflows._function_executor": ("FunctionExecutor", "executor"), "._workflows._functional": ( "FunctionalWorkflow", + "FunctionalWorkflowDefinition", "FunctionalWorkflowAgent", "RunContext", "StepWrapper", @@ -476,6 +477,7 @@ "FunctionTool", "FunctionalWorkflow", "FunctionalWorkflowAgent", + "FunctionalWorkflowDefinition", "GeneratedEmbeddings", "GraphConnectivityError", "HistoryProvider", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 31fa53a56ff..f4052e798c0 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -283,6 +283,7 @@ from ._workflows._function_executor import FunctionExecutor, executor from ._workflows._functional import ( FunctionalWorkflow, FunctionalWorkflowAgent, + FunctionalWorkflowDefinition, RunContext, StepWrapper, get_run_context, @@ -440,6 +441,7 @@ __all__ = [ "FunctionTool", "FunctionalWorkflow", "FunctionalWorkflowAgent", + "FunctionalWorkflowDefinition", "GeneratedEmbeddings", "GraphConnectivityError", "HistoryProvider", diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 7c38ef1acf4..9e4deb340bb 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -21,7 +21,10 @@ Key public symbols: -* :func:`workflow` / :class:`FunctionalWorkflow` — decorator and runtime. +* :func:`workflow` / :class:`FunctionalWorkflowDefinition` — decorator and + stateless definition. +* :class:`FunctionalWorkflow` — stateful runtime created by + :meth:`FunctionalWorkflowDefinition.build`. * :func:`step` / :class:`StepWrapper` — optional step decorator. * :class:`RunContext` — execution context injected into workflow and step functions. @@ -628,6 +631,46 @@ def _decorator(fn: Callable[..., Awaitable[Any]]) -> StepWrapper[Any]: return _decorator +# --------------------------------------------------------------------------- +# FunctionalWorkflowDefinition +# --------------------------------------------------------------------------- + + +@experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) +class FunctionalWorkflowDefinition: + """Stateless definition produced by :func:`workflow`. + + Call :meth:`build` to create a stateful :class:`FunctionalWorkflow`. + Each built workflow represents one logical caller or session. + """ + + def __init__( + self, + func: Callable[..., Awaitable[Any]], + *, + name: str | None = None, + description: str | None = None, + ) -> None: + FunctionalWorkflow._classify_signature(func) + self._func = func + self.name = name or func.__name__ + self.description = description + functools.update_wrapper(self, func) # type: ignore[arg-type] + + def build( + self, + *, + checkpoint_storage: CheckpointStorage | None = None, + ) -> FunctionalWorkflow: + """Build a stateful workflow for one logical caller or session.""" + return FunctionalWorkflow( + self._func, + name=self.name, + description=self.description, + checkpoint_storage=checkpoint_storage, + ) + + # --------------------------------------------------------------------------- # FunctionalWorkflow # --------------------------------------------------------------------------- @@ -637,8 +680,8 @@ def _decorator(fn: Callable[..., Awaitable[Any]]) -> StepWrapper[Any]: class FunctionalWorkflow: """A workflow backed by a user-defined async function. - Created by the :func:`workflow` decorator. Exposes the same ``run()`` - interface as graph-based :class:`Workflow` objects, returning a + Built from a :class:`FunctionalWorkflowDefinition`. Exposes the same + ``run()`` interface as graph-based :class:`Workflow` objects, returning a :class:`WorkflowRunResult` (or a :class:`ResponseStream` in streaming mode). @@ -648,11 +691,7 @@ class FunctionalWorkflow: Like graph-based :class:`Workflow`, each instance owns mutable execution state across calls to :meth:`run`. Scope an instance to one logical - caller or session. Use :meth:`create_instance` when independent callers - need the same workflow definition; do not share one instance across - mutually untrusted callers. Instance isolation covers in-memory state. - When checkpointing is enabled, the host must also authorize and - tenant-scope access to the checkpoint storage. + caller or session; build separate instances for independent callers. Args: func: The async function that implements the workflow logic. @@ -672,7 +711,8 @@ async def my_pipeline(data: str) -> str: return await to_upper(data) - result = await my_pipeline.run("hello") + pipeline = my_pipeline.build() + result = await pipeline.run("hello") print(result.get_outputs()) # ['HELLO'] """ @@ -707,32 +747,6 @@ def __init__( functools.update_wrapper(self, func) # type: ignore[arg-type] - def create_instance( - self, - *, - checkpoint_storage: CheckpointStorage | None = None, - ) -> FunctionalWorkflow: - """Create an independent stateful instance of this workflow definition. - - The new instance reuses the workflow function and configuration but - has its own message, step-cache, and pending-request state. Checkpoint - storage is not inherited because a shared external store requires - host-managed authorization and tenant isolation. - - Args: - checkpoint_storage: Optional storage scoped and authorized for the - new instance's logical caller or session. - - Returns: - A new :class:`FunctionalWorkflow` for one logical caller or session. - """ - return FunctionalWorkflow( - self._func, - name=self.name, - description=self.description, - checkpoint_storage=checkpoint_storage, - ) - @staticmethod def _classify_signature(func: Callable[..., Any]) -> list[str]: """Return the names of non-ctx parameters, validating arity. @@ -967,7 +981,7 @@ async def _run_core( if storage is None: raise ValueError( "Cannot restore from checkpoint without checkpoint_storage. " - "Provide checkpoint_storage parameter or set it on the @workflow decorator." + "Provide checkpoint_storage to build() or to this run." ) checkpoint = await storage.load(checkpoint_id) if checkpoint.graph_signature_hash != self.graph_signature_hash: @@ -1292,7 +1306,7 @@ async def _run_cleanup(self) -> None: @overload -def workflow(func: Callable[..., Awaitable[Any]]) -> FunctionalWorkflow: ... +def workflow(func: Callable[..., Awaitable[Any]]) -> FunctionalWorkflowDefinition: ... @overload @@ -1300,8 +1314,7 @@ def workflow( *, name: str | None = None, description: str | None = None, - checkpoint_storage: CheckpointStorage | None = None, -) -> Callable[[Callable[..., Awaitable[Any]]], FunctionalWorkflow]: ... +) -> Callable[[Callable[..., Awaitable[Any]]], FunctionalWorkflowDefinition]: ... @experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) @@ -1310,29 +1323,26 @@ def workflow( *, name: str | None = None, description: str | None = None, - checkpoint_storage: CheckpointStorage | None = None, -) -> FunctionalWorkflow | Callable[[Callable[..., Awaitable[Any]]], FunctionalWorkflow]: - """Decorator that converts an async function into a :class:`FunctionalWorkflow`. +) -> FunctionalWorkflowDefinition | Callable[[Callable[..., Awaitable[Any]]], FunctionalWorkflowDefinition]: + """Decorator that creates a stateless :class:`FunctionalWorkflowDefinition`. Supports both bare ``@workflow`` and parameterized ``@workflow(name="my_wf")`` forms. The decorated function receives its input as the first positional argument and a :class:`RunContext` instance wherever a parameter is annotated with - that type. The resulting :class:`FunctionalWorkflow` object exposes the - same ``run()`` interface as graph-based workflows. + that type. Call ``build()`` on the resulting definition to create a + stateful :class:`FunctionalWorkflow`. Args: func: The async function to decorate (when using the bare ``@workflow`` form). name: Display name for the workflow. Defaults to ``func.__name__``. description: Optional human-readable description. - checkpoint_storage: Default :class:`CheckpointStorage` for - persisting step results and workflow state. Returns: - A :class:`FunctionalWorkflow` (bare form) or a decorator that - produces one (parameterized form). + A :class:`FunctionalWorkflowDefinition` (bare form) or a decorator + that produces one (parameterized form). Examples: @@ -1345,14 +1355,17 @@ async def pipeline(data: str) -> str: # Parameterized form - @workflow(name="my_pipeline", checkpoint_storage=storage) + @workflow(name="my_pipeline") async def pipeline(data: str) -> str: ... + + + instance = pipeline.build(checkpoint_storage=storage) """ if func is not None: - return FunctionalWorkflow(func, name=name, description=description, checkpoint_storage=checkpoint_storage) + return FunctionalWorkflowDefinition(func, name=name, description=description) - def _decorator(fn: Callable[..., Awaitable[Any]]) -> FunctionalWorkflow: - return FunctionalWorkflow(fn, name=name, description=description, checkpoint_storage=checkpoint_storage) + def _decorator(fn: Callable[..., Awaitable[Any]]) -> FunctionalWorkflowDefinition: + return FunctionalWorkflowDefinition(fn, name=name, description=description) return _decorator diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 1ddfc1954a7..96b7e6edf69 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -7,17 +7,20 @@ import asyncio import json import logging -from collections.abc import Iterator +from collections.abc import Awaitable, Callable, Iterator from contextlib import contextmanager from dataclasses import dataclass +from typing import Any, overload import pytest from agent_framework import ( AgentResponseUpdate, + CheckpointStorage, ExperimentalFeature, FunctionalWorkflow, FunctionalWorkflowAgent, + FunctionalWorkflowDefinition, InMemoryCheckpointStorage, RunContext, StepWrapper, @@ -37,6 +40,34 @@ # --------------------------------------------------------------------------- +@overload +def built_workflow(func: Callable[..., Awaitable[Any]]) -> FunctionalWorkflow: ... + + +@overload +def built_workflow( + *, + name: str | None = None, + description: str | None = None, + checkpoint_storage: CheckpointStorage | None = None, +) -> Callable[[Callable[..., Awaitable[Any]]], FunctionalWorkflow]: ... + + +def built_workflow( + func: Callable[..., Awaitable[Any]] | None = None, + *, + name: str | None = None, + description: str | None = None, + checkpoint_storage: CheckpointStorage | None = None, +) -> FunctionalWorkflow | Callable[[Callable[..., Awaitable[Any]]], FunctionalWorkflow]: + """Build a fresh executable workflow for behavior-focused tests.""" + + def decorate(fn: Callable[..., Awaitable[Any]]) -> FunctionalWorkflow: + return workflow(name=name, description=description)(fn).build(checkpoint_storage=checkpoint_storage) + + return decorate(func) if func is not None else decorate + + @step async def add_one(x: int) -> int: return x + 1 @@ -69,7 +100,7 @@ async def failing_step(x: int) -> int: class TestBasicExecution: async def test_simple_sequential_pipeline(self): - @workflow + @built_workflow async def pipeline(x: int) -> int: a = await add_one(x) return await double(a) @@ -80,7 +111,7 @@ async def pipeline(x: int) -> int: assert outputs == [12] # (5+1)*2 async def test_workflow_with_string_data(self): - @workflow + @built_workflow async def upper_pipeline(text: str) -> str: return await to_upper(text) @@ -88,7 +119,7 @@ async def upper_pipeline(text: str) -> str: assert result.get_outputs() == ["HELLO"] async def test_workflow_returns_result(self): - @workflow + @built_workflow async def simple(x: int) -> int: return await add_one(x) @@ -96,14 +127,14 @@ async def simple(x: int) -> int: assert result.get_outputs() == [11] async def test_workflow_name_defaults_to_function_name(self): - @workflow + @built_workflow async def my_pipeline(x: int) -> int: return x assert my_pipeline.name == "my_pipeline" async def test_workflow_custom_name(self): - @workflow(name="custom_wf", description="A test workflow") + @built_workflow(name="custom_wf", description="A test workflow") async def wf(x: int) -> int: return x @@ -118,7 +149,7 @@ async def wf(x: int) -> int: class TestEventEmission: async def test_step_events_emitted(self): - @workflow + @built_workflow async def pipeline(x: int) -> int: return await add_one(x) @@ -129,7 +160,7 @@ async def pipeline(x: int) -> int: assert "output" in event_types async def test_step_events_carry_executor_id(self): - @workflow + @built_workflow async def pipeline(x: int) -> int: return await add_one(x) @@ -144,7 +175,7 @@ async def pipeline(x: int) -> int: assert completed_events[0].data == 6 async def test_status_events_in_timeline(self): - @workflow + @built_workflow async def pipeline(x: int) -> int: return x @@ -154,7 +185,7 @@ async def pipeline(x: int) -> int: assert WorkflowRunState.IDLE in states async def test_final_state_is_idle(self): - @workflow + @built_workflow async def pipeline(x: int) -> int: return x @@ -164,7 +195,7 @@ async def pipeline(x: int) -> int: async def test_custom_event(self): from agent_framework import WorkflowEvent - @workflow + @built_workflow async def pipeline(x: int, ctx: RunContext) -> int: await ctx.add_event(WorkflowEvent("intermediate", executor_id="pipeline", data="custom_data")) return x @@ -192,7 +223,7 @@ async def slow_double(x: int) -> int: await asyncio.sleep(0.01) return x * 2 - @workflow + @built_workflow async def parallel_wf(x: int) -> list[int]: a, b = await asyncio.gather(slow_add(x), slow_double(x)) return [a, b] @@ -210,7 +241,7 @@ async def task_a(x: int) -> int: async def task_b(x: int) -> int: return x * 2 - @workflow + @built_workflow async def par_wf(x: int) -> tuple[int, int]: a, b = await asyncio.gather(task_a(x), task_b(x)) return (a, b) @@ -228,14 +259,16 @@ async def par_wf(x: int) -> tuple[int, int]: class TestHITL: - async def test_separate_instances_isolate_pending_continuations(self): + async def test_workflow_definition_builds_isolated_pending_continuations(self): @workflow async def review_wf(doc: str, ctx: RunContext) -> str: feedback = await ctx.request_info(doc, response_type=str) return f"{doc}:{feedback}" - caller_a = review_wf.create_instance() - caller_b = review_wf.create_instance() + assert not hasattr(review_wf, "run") + + caller_a = review_wf.build() + caller_b = review_wf.build() caller_a_paused = await caller_a.run("caller-a") request_id = caller_a_paused.get_request_info_events()[0].request_id @@ -246,34 +279,30 @@ async def review_wf(doc: str, ctx: RunContext) -> str: caller_a_completed = await caller_a.run(responses={request_id: "caller-a-response"}) assert caller_a_completed.get_outputs() == ["caller-a:caller-a-response"] - async def test_create_instance_does_not_inherit_checkpoint_storage_by_default(self): - storage = InMemoryCheckpointStorage() - - @workflow(checkpoint_storage=storage) - async def review_wf(doc: str, ctx: RunContext) -> str: - return await ctx.request_info(doc, response_type=str) + async def test_build_does_not_inherit_checkpoint_storage_by_default(self): + @workflow + async def review_wf(doc: str) -> str: + return doc - caller = review_wf.create_instance() - await caller.run("caller") + caller = review_wf.build() - assert await storage.list_checkpoints(workflow_name="review_wf") == [] + with pytest.raises(ValueError, match="checkpoint_storage"): + await caller.run(checkpoint_id="missing") - async def test_create_instance_accepts_tenant_scoped_checkpoint_storage(self): - definition_storage = InMemoryCheckpointStorage() + async def test_build_accepts_tenant_scoped_checkpoint_storage(self): caller_storage = InMemoryCheckpointStorage() - @workflow(checkpoint_storage=definition_storage) + @workflow async def review_wf(doc: str, ctx: RunContext) -> str: return await ctx.request_info(doc, response_type=str) - caller = review_wf.create_instance(checkpoint_storage=caller_storage) + caller = review_wf.build(checkpoint_storage=caller_storage) await caller.run("caller") - assert await definition_storage.list_checkpoints(workflow_name="review_wf") == [] assert len(await caller_storage.list_checkpoints(workflow_name="review_wf")) == 1 async def test_request_info_interrupts(self): - @workflow + @built_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}" @@ -286,7 +315,7 @@ async def review_wf(doc: str, ctx: RunContext) -> str: assert request_events[0].request_id == "req1" async def test_request_info_resume(self): - @workflow + @built_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}" @@ -304,7 +333,7 @@ async def review_wf(doc: str, ctx: RunContext) -> str: 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.""" - @workflow + @built_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}" @@ -323,7 +352,7 @@ async def review_wf(doc: str, ctx: RunContext) -> str: 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.""" - @workflow + @built_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}" @@ -341,7 +370,7 @@ async def review_wf(doc: str, ctx: RunContext) -> str: async def test_untyped_ctx_parameter(self): """ctx is injected by parameter name even without a RunContext annotation.""" - @workflow # pyright: ignore[reportUnknownArgumentType] + @built_workflow # pyright: ignore[reportUnknownArgumentType] async def review_wf(doc: str, ctx) -> str: # pyright: ignore[reportMissingParameterType, reportUnknownParameterType] feedback: str = await ctx.request_info({"draft": doc}, response_type=str, request_id="req1") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] return f"Final: {feedback}" @@ -353,7 +382,7 @@ async def review_wf(doc: str, ctx) -> str: # pyright: ignore[reportMissingParam assert result2.get_outputs() == ["Final: LGTM"] async def test_multiple_sequential_interrupts(self): - @workflow + @built_workflow 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") @@ -374,7 +403,7 @@ async def multi_hitl(data: str, ctx: RunContext) -> str: assert result3.get_outputs() == ["A+B"] async def test_request_info_auto_generates_id(self): - @workflow + @built_workflow async def auto_id_wf(x: int, ctx: RunContext) -> None: await ctx.request_info("need data", response_type=str) @@ -391,7 +420,7 @@ async def auto_id_wf(x: int, ctx: RunContext) -> None: class TestErrorHandling: async def test_step_failure_propagates(self): - @workflow + @built_workflow async def failing_wf(x: int) -> None: await failing_step(x) @@ -399,7 +428,7 @@ async def failing_wf(x: int) -> None: await failing_wf.run(42) async def test_step_failure_emits_executor_failed(self): - @workflow + @built_workflow async def failing_wf(x: int) -> None: await failing_step(x) @@ -415,7 +444,7 @@ async def failing_wf(x: int) -> None: assert failed_events[0].executor_id == "failing_step" async def test_workflow_failure_emits_failed_status(self): - @workflow + @built_workflow async def bad_wf(x: int) -> None: raise RuntimeError("workflow broke") @@ -431,7 +460,7 @@ async def bad_wf(x: int) -> None: assert any(e.state == WorkflowRunState.FAILED for e in status_events) async def test_invalid_params_message_and_responses(self): - @workflow + @built_workflow async def wf(x: int) -> None: pass @@ -439,7 +468,7 @@ async def wf(x: int) -> None: await wf.run("hello", responses={"r1": "val"}) async def test_invalid_params_message_and_checkpoint(self): - @workflow + @built_workflow async def wf(x: int) -> None: pass @@ -447,7 +476,7 @@ async def wf(x: int) -> None: await wf.run("hello", checkpoint_id="abc") async def test_invalid_params_nothing(self): - @workflow + @built_workflow async def wf(x: int) -> None: pass @@ -462,7 +491,7 @@ async def wf(x: int) -> None: class TestStreaming: async def test_streaming_yields_events(self): - @workflow + @built_workflow async def pipeline(x: int) -> int: return await add_one(x) @@ -478,7 +507,7 @@ async def pipeline(x: int) -> int: assert "output" in event_types async def test_streaming_final_response(self): - @workflow + @built_workflow async def pipeline(x: int) -> int: return await add_one(x) @@ -490,7 +519,7 @@ async def pipeline(x: int) -> int: async def test_streaming_context_reports_streaming(self): streaming_flag = None - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> int: nonlocal streaming_flag streaming_flag = ctx.is_streaming() # type: ignore[assignment] @@ -535,7 +564,7 @@ def test_step_wrapper_is_step_wrapper(self): class TestStateManagement: async def test_get_set_state(self): - @workflow + @built_workflow async def stateful_wf(x: int, ctx: RunContext) -> int: ctx.set_state("counter", x) return ctx.get_state("counter") @@ -544,7 +573,7 @@ async def stateful_wf(x: int, ctx: RunContext) -> int: assert result.get_outputs() == [42] async def test_get_state_default(self): - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> str: return ctx.get_state("missing", "default_val") @@ -565,7 +594,7 @@ async def test_checkpoint_save_and_restore(self): async def expensive(x: int) -> int: return x * 100 - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def ckpt_wf(x: int) -> int: return await expensive(x) @@ -583,7 +612,7 @@ async def test_checkpoint_runtime_storage_override(self): async def compute(x: int) -> int: return x + 1 - @workflow + @built_workflow async def wf(x: int) -> int: return await compute(x) @@ -603,7 +632,7 @@ async def counting_task(x: int) -> int: call_count += 1 return x + 1 - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def wf(x: int) -> int: return await counting_task(x) @@ -624,7 +653,7 @@ async def wf(x: int) -> int: async def test_checkpoint_hitl_resume(self): storage = InMemoryCheckpointStorage() - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def hitl_wf(doc: str, ctx: RunContext) -> str: feedback = await ctx.request_info({"draft": doc}, response_type=str, request_id="req1") return f"Done: {feedback}" @@ -642,7 +671,7 @@ async def hitl_wf(doc: str, ctx: RunContext) -> str: assert result2.get_outputs() == ["Done: Approved!"] async def test_checkpoint_without_storage_raises(self): - @workflow + @built_workflow async def wf(x: int) -> int: return x @@ -652,7 +681,7 @@ async def wf(x: int) -> int: async def test_checkpoint_preserves_state(self): storage = InMemoryCheckpointStorage() - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def stateful_wf(x: int, ctx: RunContext) -> str: ctx.set_state("key", "value") feedback = await ctx.request_info("need info", response_type=str, request_id="r1") @@ -692,7 +721,7 @@ async def crashing_step2(x: int) -> int: raise RuntimeError("simulated crash") return x * 2 - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def crash_wf(x: int) -> int: a = await slow_step1(x) return await crashing_step2(a) @@ -731,7 +760,7 @@ async def s2(x: int) -> int: async def s3(x: int) -> int: return x + 3 - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def multi_step_wf(x: int) -> int: a = await s1(x) b = await s2(a) @@ -752,7 +781,7 @@ async def test_no_checkpoint_on_cache_hit(self): async def compute(x: int) -> int: return x + 1 - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def wf(x: int) -> int: return await compute(x) @@ -792,7 +821,7 @@ async def process_normal(text: str) -> str: async def quarantine(text: str) -> str: return f"quarantined: {text}" - @workflow + @built_workflow async def email_pipeline(email: str) -> str: cl = await classify(email) if cl.is_spam: @@ -819,7 +848,7 @@ async def test_nested_workflow_as_task(self): async def step_a(x: int) -> int: return x + 1 - @workflow + @built_workflow async def inner_wf(x: int) -> int: return await step_a(x) @@ -828,7 +857,7 @@ async def call_inner(x: int) -> int: result = await inner_wf.run(x) return result.get_outputs()[0] - @workflow + @built_workflow async def outer_wf(x: int) -> int: return await call_inner(x) @@ -843,7 +872,7 @@ async def outer_wf(x: int) -> int: class TestAsAgent: async def test_as_agent_returns_agent(self): - @workflow + @built_workflow async def wf(x: int) -> str: return f"result: {x}" @@ -851,7 +880,7 @@ async def wf(x: int) -> str: assert agent.name == "wf" async def test_as_agent_custom_name(self): - @workflow + @built_workflow async def wf(x: int) -> int: return x @@ -859,7 +888,7 @@ async def wf(x: int) -> int: assert agent.name == "my_agent" async def test_as_agent_run(self): - @workflow + @built_workflow async def wf(x: int) -> int: return await add_one(x) @@ -868,7 +897,7 @@ async def wf(x: int) -> int: assert response.text == "11" async def test_as_agent_run_streaming(self): - @workflow + @built_workflow async def wf(x: int) -> str: return f"result: {x}" @@ -884,7 +913,7 @@ async def wf(x: int) -> str: assert len(response.messages) >= 1 async def test_as_agent_has_id_and_description(self): - @workflow(description="A test workflow") + @built_workflow(description="A test workflow") async def wf(x: int) -> int: return x @@ -900,7 +929,7 @@ async def wf(x: int) -> int: class TestConcurrencyGuard: async def test_concurrent_run_raises(self): - @workflow + @built_workflow async def slow_wf(x: int) -> int: await asyncio.sleep(0.1) return x @@ -916,7 +945,7 @@ async def slow_wf(x: int) -> int: await stream.get_final_response() async def test_run_after_completion(self): - @workflow + @built_workflow async def wf(x: int) -> int: return x @@ -955,7 +984,7 @@ def test_workflow_bare_decorator(self): async def my_wf(x: int) -> None: pass - assert isinstance(my_wf, FunctionalWorkflow) + assert isinstance(my_wf, FunctionalWorkflowDefinition) assert my_wf.name == "my_wf" def test_workflow_with_params(self): @@ -963,7 +992,7 @@ def test_workflow_with_params(self): async def my_wf(x: int) -> None: pass - assert isinstance(my_wf, FunctionalWorkflow) + assert isinstance(my_wf, FunctionalWorkflowDefinition) assert my_wf.name == "custom" assert my_wf.description == "desc" @@ -975,7 +1004,7 @@ async def my_wf(x: int) -> None: class TestIncludeStatusEvents: async def test_status_events_excluded_by_default(self): - @workflow + @built_workflow async def wf(x: int) -> int: return x @@ -984,7 +1013,7 @@ async def wf(x: int) -> int: assert len(status_in_list) == 0 async def test_status_events_included_when_requested(self): - @workflow + @built_workflow async def wf(x: int) -> int: return x @@ -1000,7 +1029,7 @@ async def wf(x: int) -> int: class TestEdgeCases: async def test_workflow_with_no_tasks(self): - @workflow + @built_workflow async def no_tasks(x: int) -> int: return x * 2 @@ -1008,7 +1037,7 @@ async def no_tasks(x: int) -> int: assert result.get_outputs() == [10] async def test_workflow_with_no_output(self): - @workflow + @built_workflow async def silent_wf(x: int) -> None: pass # returns None — no output emitted @@ -1018,7 +1047,7 @@ async def silent_wf(x: int) -> None: async def test_return_value_auto_yields_output(self): """Returning a non-None value automatically emits it as an output.""" - @workflow + @built_workflow async def wf(x: int) -> int: return x * 3 @@ -1026,7 +1055,7 @@ async def wf(x: int) -> int: assert result.get_outputs() == [15] async def test_step_called_multiple_times(self): - @workflow + @built_workflow async def wf(x: int) -> int: a = await add_one(x) b = await add_one(a) @@ -1049,7 +1078,7 @@ async def wf(x: int) -> int: class TestRecoveryAfterErrors: async def test_run_after_failure_is_allowed(self): - @workflow + @built_workflow async def wf(x: int) -> int: if x == 1: raise RuntimeError("boom") @@ -1080,7 +1109,7 @@ async def test_except_exception_does_not_catch_interrupt(self): """User code with ``except Exception`` should not catch WorkflowInterrupted.""" caught = False - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> str: nonlocal caught try: @@ -1108,7 +1137,7 @@ async def test_checkpoint_signature_mismatch_raises(self): storage = InMemoryCheckpointStorage() - @workflow(name="my_wf", checkpoint_storage=storage) + @built_workflow(name="my_wf", checkpoint_storage=storage) async def wf(x: int) -> int: return x @@ -1152,7 +1181,7 @@ async def tracked(x: int) -> int: call_count += 1 return x + 1 - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def wf(x: int) -> int: return await tracked(x) @@ -1185,7 +1214,7 @@ async def test_bypassed_event_carries_cached_data(self): async def compute(x: int) -> int: return x * 10 - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def wf(x: int) -> int: return await compute(x) @@ -1213,7 +1242,7 @@ async def review_step(doc: str, ctx: RunContext) -> str: feedback = await ctx.request_info({"draft": doc}, response_type=str, request_id="s1") return f"reviewed: {feedback}" - @workflow + @built_workflow async def wf(doc: str) -> str: return await review_step(doc) @@ -1249,7 +1278,7 @@ async def needs_ctx_first(ctx: RunContext, data: str) -> str: ctx.set_state("seen", data) return f"{data}:{ctx.get_state('seen')}" - @workflow + @built_workflow async def wf(data: str) -> str: return await needs_ctx_first(data) @@ -1269,7 +1298,7 @@ async def capture_ctx(x: int) -> int: captured_ctx = get_run_context() # type: ignore[assignment] return x - @workflow + @built_workflow async def wf(x: int) -> int: return await capture_ctx(x) @@ -1293,7 +1322,7 @@ class TestNoneResponseHandling: async def test_none_response_logs_warning(self): """Providing None as a response value should log a warning.""" - @workflow + @built_workflow async def wf(doc: str, ctx: RunContext) -> str: val = await ctx.request_info("need input", response_type=str, request_id="r1") return f"got: {val}" @@ -1311,7 +1340,7 @@ async def wf(doc: str, ctx: RunContext) -> str: async def test_none_response_is_returned(self): """None is a valid (if discouraged) response value.""" - @workflow + @built_workflow 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}" @@ -1366,7 +1395,7 @@ async def step_b(val: int, ctx: RunContext) -> str: feedback = await ctx.request_info({"val": val}, response_type=str, request_id="r1") return f"{val}:{feedback}" - @workflow + @built_workflow async def wf(x: int) -> str: a = await step_a(x) return await step_b(a) @@ -1397,7 +1426,7 @@ async def review(val: int, ctx: RunContext) -> str: feedback = await ctx.request_info({"val": val}, response_type=str, request_id="rev") return f"reviewed({val}):{feedback}" - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def wf(x: int) -> str: v = await compute(x) return await review(v) @@ -1428,7 +1457,7 @@ async def needs_feedback(doc: str, ctx: RunContext) -> str: val = await ctx.request_info({"doc": doc}, response_type=str, request_id="r1") return f"got:{val}" - @workflow + @built_workflow async def wf(doc: str) -> str: return await needs_feedback(doc) @@ -1447,7 +1476,7 @@ async def test_step_hitl_does_not_emit_executor_failed(self): async def hitl_step(x: int, ctx: RunContext) -> str: return await ctx.request_info("need data", response_type=str, request_id="r1") - @workflow + @built_workflow async def wf(x: int) -> str: return await hitl_step(x) @@ -1466,7 +1495,7 @@ class TestDeterministicAutoRequestId: """Regression for bug_001: auto-generated request_info ids must be stable across replay.""" async def test_auto_request_id_roundtrips_on_resume(self): - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> str: # No request_id — framework must generate a deterministic one val = await ctx.request_info("need data", response_type=str) @@ -1485,7 +1514,7 @@ async def wf(x: int, ctx: RunContext) -> str: assert result2.get_outputs() == ["got:hello"] async def test_multiple_auto_ids_are_distinct_and_stable(self): - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> str: a = await ctx.request_info("first", response_type=str) b = await ctx.request_info("second", response_type=str) @@ -1512,7 +1541,7 @@ async def first_review(value: int, ctx: RunContext) -> str: async def second_review(value: int, ctx: RunContext) -> str: return await ctx.request_info({"step": "second", "value": value}, response_type=str) - @workflow + @built_workflow async def wf(value: int) -> str: first = await first_review(value) second = await second_review(value) @@ -1539,7 +1568,7 @@ class TestPendingRequestsPruned: async def test_final_checkpoint_no_longer_claims_resolved_requests_pending(self): storage = InMemoryCheckpointStorage() - @workflow(checkpoint_storage=storage) + @built_workflow(checkpoint_storage=storage) async def wf(x: int, ctx: RunContext) -> str: a = await ctx.request_info("q1", response_type=str, request_id="r1") b = await ctx.request_info("q2", response_type=str, request_id="r2") @@ -1567,7 +1596,7 @@ async def wf(a: str, b: str, ctx: RunContext) -> str: return f"{a}+{b}" async def test_ctx_only_workflow_with_message_raises_clear_error(self): - @workflow + @built_workflow async def wf(ctx: RunContext) -> str: return "no message used" @@ -1579,7 +1608,7 @@ def test_ctx_only_workflow_decoration_succeeds(self): # message-receiving parameter. (Running it without a message still # requires providing responses or a checkpoint_id — that's # _validate_run_params's job, not ours.) - @workflow + @built_workflow async def wf(ctx: RunContext) -> str: return "ok" @@ -1590,7 +1619,7 @@ class TestStaleResponsesRejected: """Regression for bug_014: stale responses after clean completion must be rejected.""" async def test_responses_after_clean_completion_raise(self): - @workflow + @built_workflow async def wf(x: int) -> int: return x * 2 @@ -1599,7 +1628,7 @@ async def wf(x: int) -> int: await wf.run(responses={"stale": "x"}) async def test_responses_mismatched_key_raises(self): - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> str: return await ctx.request_info("q", response_type=str, request_id="r1") @@ -1612,7 +1641,7 @@ class TestReservedStateKeys: """Regression for bug_017: set_state must reject underscore-prefixed keys.""" async def test_underscore_key_rejected(self): - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> int: ctx.set_state("_private", "user value") return x @@ -1621,7 +1650,7 @@ async def wf(x: int, ctx: RunContext) -> int: await wf.run(1) async def test_normal_key_still_works(self): - @workflow + @built_workflow async def wf(x: int, ctx: RunContext) -> int: ctx.set_state("normal_key", "v") assert ctx.get_state("normal_key") == "v" @@ -1641,7 +1670,7 @@ async def test_step_with_non_deepcopyable_arg_replays(self): async def takes_lock(lock: threading.Lock, n: int) -> int: return n + 1 - @workflow + @built_workflow async def wf(x: int) -> int: lock = threading.Lock() return await takes_lock(lock, x) @@ -1657,11 +1686,11 @@ class TestStepDiscoveryAttributeAccess: """Regression for bug_008: checkpoint hash must differ when function body changes.""" async def test_signature_hash_changes_when_function_body_changes(self): - @workflow + @built_workflow async def wf_a(x: int) -> int: return x + 1 - @workflow(name="wf_b") + @built_workflow(name="wf_b") async def wf_b(x: int) -> int: return x * 100 @@ -1674,7 +1703,7 @@ class TestAsAgentSignatureParity: """Regression for bug_015: as_agent signature must accept description/context_providers.""" async def test_as_agent_accepts_description_override(self): - @workflow(description="workflow level") + @built_workflow(description="workflow level") async def wf(x: str) -> str: return x.upper() @@ -1682,7 +1711,7 @@ async def wf(x: str) -> str: assert agent.description == "agent level" async def test_as_agent_accepts_context_providers_kwarg(self): - @workflow + @built_workflow async def wf(x: str) -> str: return x @@ -1691,7 +1720,7 @@ async def wf(x: str) -> str: assert list(agent.context_providers or []) == providers async def test_as_agent_description_defaults_to_workflow_description(self): - @workflow(description="from workflow") + @built_workflow(description="from workflow") async def wf(x: str) -> str: return x @@ -1703,7 +1732,7 @@ class TestFunctionalWorkflowAgentHITL: """Regression for bug_013: .as_agent() must surface request_info events.""" async def test_request_info_surfaces_as_function_approval_request(self): - @workflow + @built_workflow async def wf(x: str, ctx: RunContext) -> str: answer = await ctx.request_info({"need": x}, response_type=str, request_id="rid-1") return f"got:{answer}" @@ -1730,7 +1759,7 @@ class HandoffRequest: target_agent: str reason: str - @workflow + @built_workflow async def wf(x: str, ctx: RunContext) -> str: answer = await ctx.request_info( HandoffRequest(target_agent=x, reason="overflow"), @@ -1756,7 +1785,7 @@ async def wf(x: str, ctx: RunContext) -> str: assert json.loads(json.dumps(function_call_arguments)) == function_call_arguments async def test_resume_via_agent_responses_kwarg(self): - @workflow + @built_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}" @@ -1795,6 +1824,7 @@ def test_public_symbols_are_marked_experimental(self) -> None: StepWrapper, step, FunctionalWorkflow, + FunctionalWorkflowDefinition, workflow, FunctionalWorkflowAgent, ] diff --git a/python/samples/01-get-started/05_functional_workflow_with_agents.py b/python/samples/01-get-started/05_functional_workflow_with_agents.py index 9aff97826d1..e77c1bc7431 100644 --- a/python/samples/01-get-started/05_functional_workflow_with_agents.py +++ b/python/samples/01-get-started/05_functional_workflow_with_agents.py @@ -47,7 +47,8 @@ async def poem_workflow(topic: str) -> str: async def main() -> None: - result = await poem_workflow.run("a cat learning to code") + workflow_instance = poem_workflow.build() + result = await workflow_instance.run("a cat learning to code") print(result.get_outputs()[0]) diff --git a/python/samples/01-get-started/06_functional_workflow_basics.py b/python/samples/01-get-started/06_functional_workflow_basics.py index a679e70a995..ec6fd79f538 100644 --- a/python/samples/01-get-started/06_functional_workflow_basics.py +++ b/python/samples/01-get-started/06_functional_workflow_basics.py @@ -43,7 +43,8 @@ async def text_workflow(text: str) -> str: async def main() -> None: # - result = await text_workflow.run("hello world") + workflow_instance = text_workflow.build() + result = await workflow_instance.run("hello world") print(f"Output: {result.get_outputs()}") print(f"Final state: {result.get_final_state()}") # diff --git a/python/samples/03-workflows/functional/agent_integration.py b/python/samples/03-workflows/functional/agent_integration.py index b5911cb690f..23699e271f4 100644 --- a/python/samples/03-workflows/functional/agent_integration.py +++ b/python/samples/03-workflows/functional/agent_integration.py @@ -94,12 +94,15 @@ async def cached_pipeline(document: str) -> str: async def main(): + simple_workflow = simple_pipeline.build() + cached_workflow = cached_pipeline.build() + # Simple version — agents called inline - result = await simple_pipeline.run("This is a technical document about machine learning...") + result = await simple_workflow.run("This is a technical document about machine learning...") print(result.get_outputs()[0]) # Cached version — same result, but steps won't re-execute on resume - result = await cached_pipeline.run("This is a technical document about machine learning...") + result = await cached_workflow.run("This is a technical document about machine learning...") print(f"\nCached: {result.get_outputs()[0]}") diff --git a/python/samples/03-workflows/functional/basic_pipeline.py b/python/samples/03-workflows/functional/basic_pipeline.py index 81514da53a4..eb32b54c671 100644 --- a/python/samples/03-workflows/functional/basic_pipeline.py +++ b/python/samples/03-workflows/functional/basic_pipeline.py @@ -23,8 +23,8 @@ async def transform_data(data: dict[str, str | int]) -> str: return f"[{data['status']}] {data['content']}" -# @workflow turns this async function into a FunctionalWorkflow object. -# Without it, this is just a normal async function. With it, you get: +# @workflow turns this async function into a stateless workflow definition. +# Call .build() to create a stateful FunctionalWorkflow with: # - .run() that returns a WorkflowRunResult with events and outputs # - .run(stream=True) for streaming events in real time # - .as_agent() to use this workflow anywhere an agent is expected @@ -48,8 +48,8 @@ async def data_pipeline(url: str) -> str: async def main(): - # .run() is provided by @workflow — a plain async function wouldn't have it - result = await data_pipeline.run("https://example.com/api/data") + workflow_instance = data_pipeline.build() + result = await workflow_instance.run("https://example.com/api/data") print("Output:", result.get_outputs()[0]) print("State:", result.get_final_state()) diff --git a/python/samples/03-workflows/functional/basic_streaming_pipeline.py b/python/samples/03-workflows/functional/basic_streaming_pipeline.py index 4ee61da60f8..a986b97ee0b 100644 --- a/python/samples/03-workflows/functional/basic_streaming_pipeline.py +++ b/python/samples/03-workflows/functional/basic_streaming_pipeline.py @@ -26,9 +26,9 @@ async def validate_result(summary: str) -> bool: return len(summary) > 0 and "[200]" in summary -# @workflow enables .run(stream=True), which returns a ResponseStream -# you can iterate over with `async for`. Without @workflow, you'd just -# have a normal async function with no streaming capability. +# @workflow creates a definition. Its built workflow enables +# .run(stream=True), which returns a ResponseStream you can iterate over with +# `async for`. @workflow async def data_pipeline(url: str) -> str: """A simple sequential data pipeline.""" @@ -40,10 +40,12 @@ async def data_pipeline(url: str) -> str: async def main(): + workflow_instance = data_pipeline.build() + # run(stream=True) returns a ResponseStream that yields events as they # are produced. The raw stream includes lifecycle events (started, status) # alongside application events — filter by event.type to find what you need. - stream = data_pipeline.run("https://example.com/api/data", stream=True) + stream = workflow_instance.run("https://example.com/api/data", stream=True) async for event in stream: if event.type == "output": print(f"Output: {event.data}") diff --git a/python/samples/03-workflows/functional/hitl_review.py b/python/samples/03-workflows/functional/hitl_review.py index acca1974a8f..a9343978a05 100644 --- a/python/samples/03-workflows/functional/hitl_review.py +++ b/python/samples/03-workflows/functional/hitl_review.py @@ -59,9 +59,11 @@ async def review_pipeline(topic: str, ctx: RunContext) -> str: async def main(): + workflow_instance = review_pipeline.build() + # Phase 1: Run until the workflow pauses for human input print("=== Phase 1: Initial run ===") - result1 = await review_pipeline.run("AI Safety") + result1 = await workflow_instance.run("AI Safety") # If request_info() was reached, the state is IDLE_WITH_PENDING_REQUESTS. # If the workflow completed without hitting request_info(), it would be IDLE. @@ -75,7 +77,7 @@ async def main(): # 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 workflow_instance.run(responses={"review_request": "Add more details about alignment research"}) print(f"State: {result2.get_final_state()}") print(f"Output: {result2.get_outputs()[0]}") diff --git a/python/samples/03-workflows/functional/naive_group_chat.py b/python/samples/03-workflows/functional/naive_group_chat.py index 45a3266049b..255b185fbf8 100644 --- a/python/samples/03-workflows/functional/naive_group_chat.py +++ b/python/samples/03-workflows/functional/naive_group_chat.py @@ -74,7 +74,8 @@ async def group_chat(question: str) -> str: async def main(): - result = await group_chat.run("What's the difference between a list and a tuple in Python?") + workflow_instance = group_chat.build() + result = await workflow_instance.run("What's the difference between a list and a tuple in Python?") print(result.get_outputs()[0]) diff --git a/python/samples/03-workflows/functional/parallel_pipeline.py b/python/samples/03-workflows/functional/parallel_pipeline.py index 21fd8dceacb..87a695bf8e5 100644 --- a/python/samples/03-workflows/functional/parallel_pipeline.py +++ b/python/samples/03-workflows/functional/parallel_pipeline.py @@ -36,9 +36,9 @@ async def synthesize(sources: list[str]) -> str: return "Research Summary:\n" + "\n".join(f" - {s}" for s in sources) -# @workflow wraps the orchestration logic so you get .run(), streaming, -# and events. The functions it calls are plain Python — no decorators -# needed just because they're inside a workflow. +# @workflow defines the orchestration logic. Build the definition to get a +# stateful workflow with .run(), streaming, and events. The functions it calls +# are plain Python — no decorators needed just because they're inside a workflow. @workflow async def research_pipeline(topic: str) -> str: """Fan-out to three research tasks, then synthesize results.""" @@ -58,7 +58,8 @@ async def research_pipeline(topic: str) -> str: async def main(): - result = await research_pipeline.run("AI agents") + workflow_instance = research_pipeline.build() + result = await workflow_instance.run("AI agents") print(result.get_outputs()[0]) diff --git a/python/samples/03-workflows/functional/steps_and_checkpointing.py b/python/samples/03-workflows/functional/steps_and_checkpointing.py index 93a9423df0a..a416491844f 100644 --- a/python/samples/03-workflows/functional/steps_and_checkpointing.py +++ b/python/samples/03-workflows/functional/steps_and_checkpointing.py @@ -55,9 +55,9 @@ async def validate_result(summary: str) -> bool: storage = InMemoryCheckpointStorage() -# checkpoint_storage tells @workflow where to persist step results. +# Build with checkpoint_storage to persist step results. # Each @step saves a checkpoint after it completes. -@workflow(checkpoint_storage=storage) +@workflow async def data_pipeline(url: str) -> str: """Mix of @step functions and plain functions.""" raw = await fetch_data(url) @@ -68,9 +68,11 @@ async def data_pipeline(url: str) -> str: async def main(): + workflow_instance = data_pipeline.build(checkpoint_storage=storage) + # --- Run 1: Everything executes normally --- print("=== Run 1: Fresh execution ===") - result = await data_pipeline.run("https://example.com/api/data") + result = await workflow_instance.run("https://example.com/api/data") print(f"Output: {result.get_outputs()[0]}") print(f"fetch_calls={fetch_calls}, transform_calls={transform_calls}") @@ -85,9 +87,10 @@ async def main(): # Only validate_result() (no @step) actually runs again. print("\n=== Run 2: Restored from checkpoint ===") latest = await storage.get_latest(workflow_name="data_pipeline") - assert latest is not None + if latest is None: + raise RuntimeError("Expected a checkpoint from the first run.") - result2 = await data_pipeline.run(checkpoint_id=latest.checkpoint_id) + result2 = await workflow_instance.run(checkpoint_id=latest.checkpoint_id) print(f"Output: {result2.get_outputs()[0]}") print(f"fetch_calls={fetch_calls}, transform_calls={transform_calls}") print("(call counts unchanged — @step results were restored from checkpoint)")