diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 8e633a9334d..cfa9aee1f84 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -199,7 +199,18 @@ 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. 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/_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 2ffe99807c0..ccfff2b0a25 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -36,10 +36,12 @@ # 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 @@ -48,7 +50,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 +65,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 +218,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 @@ -646,6 +661,29 @@ 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. + Args: func: The async function that implements the workflow logic. name: Display name for the workflow. Defaults to ``func.__name__``. @@ -687,6 +725,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 +779,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 +792,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 +805,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,9 +829,17 @@ 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). + 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), @@ -807,58 +857,48 @@ 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) - # 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. + continuation_nonce: str | None = None + if responses is not None and checkpoint_id is None: + continuation_nonce = self._validate_continuation_authority(continuation_token) 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." + ) + # 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." ) - 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), + finalizer=functools.partial( + self._finalize_events, + include_status_events=include_status_events, + continuation_token=result_continuation_token, + ), cleanup_hooks=[self._run_cleanup], ) @@ -908,6 +948,39 @@ 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 # ------------------------------------------------------------------ @@ -917,6 +990,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, @@ -963,10 +1038,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) @@ -977,7 +1048,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 @@ -996,6 +1067,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) @@ -1024,29 +1098,46 @@ 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() with _framework_event_origin(): 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._clear_continuation_state() with _framework_event_origin(): yield WorkflowEvent.status(WorkflowRunState.IDLE) span.add_event(OtelAttr.WORKFLOW_COMPLETED) except WorkflowInterrupted: - # 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) + 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() # HITL interruption — yield events collected so far for event in ctx._get_events(): @@ -1057,36 +1148,39 @@ 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: - # 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: @@ -1143,12 +1237,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, @@ -1209,6 +1304,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 +1319,56 @@ 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 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 @staticmethod def _validate_run_params( @@ -1341,7 +1486,14 @@ 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=`` 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. Args: workflow: The :class:`FunctionalWorkflow` to wrap. @@ -1379,6 +1531,16 @@ 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, @@ -1386,6 +1548,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 +1561,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 +1573,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,7 +1588,13 @@ def run( :class:`AgentResponseUpdate` items. responses: HITL responses keyed by ``request_id``, forwarded to the underlying workflow so HITL resumes work via this agent. - checkpoint_id: Optional checkpoint to restore from. + 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_storage: Override the workflow's default :class:`CheckpointStorage` for this run. **kwargs: Extra keyword arguments forwarded to the workflow run. @@ -1436,6 +1607,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 +1615,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,17 +1626,27 @@ 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: - result = await self._workflow.run( + workflow_result = 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( @@ -1471,6 +1654,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, @@ -1478,41 +1662,50 @@ 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]: - 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, - ) + 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(), @@ -1568,4 +1761,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..05f7784c67a 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,22 @@ 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) -> 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..e04e6fe7db3 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, @@ -228,6 +229,135 @@ 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_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: @@ -252,29 +382,118 @@ 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 - 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_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.""" @@ -289,7 +508,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 +527,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): @@ -313,21 +538,102 @@ 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") 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"}) - assert result3.get_outputs() == ["A+B"] + 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 async def test_request_info_auto_generates_id(self): @workflow @@ -460,6 +766,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 @@ -514,6 +839,121 @@ 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() @@ -588,6 +1028,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") @@ -618,6 +1059,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") @@ -1180,7 +1622,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 +1700,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 +1720,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 +1784,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 +1817,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 +1845,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 +1896,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 +1912,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 +1947,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 +1976,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 +2032,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 +2040,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 +2203,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 +2219,118 @@ 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.""" @@ -1740,6 +2340,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.""" diff --git a/python/samples/03-workflows/functional/hitl_review.py b/python/samples/03-workflows/functional/hitl_review.py index 39f2dae8853..08c7ea2818f 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). @@ -66,15 +66,25 @@ 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}") - - # Phase 2: Resume with the human's response + 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. 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=continuation_token, + ) print(f"State: {result2.get_final_state()}") print(f"Output: {result2.get_outputs()[0]}")