diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_hosted_workflow_conversation.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_hosted_workflow_conversation.py new file mode 100644 index 0000000000..8e3840b40e --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_hosted_workflow_conversation.py @@ -0,0 +1,193 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Hosted workflow conversation checkpoint adapter.""" + +from __future__ import annotations + +import copy +import os +from collections.abc import AsyncIterator, Callable, Sequence +from dataclasses import dataclass +from pathlib import Path + +from agent_framework import AgentResponse, AgentResponseUpdate, FileCheckpointStorage, Message, WorkflowAgent +from azure.ai.agentserver.responses.models import MessageRole + +AZURE_RESPONSES_MESSAGE_ROLE_TYPE = f"{MessageRole.__module__}:{MessageRole.__qualname__}" + + +def checkpoint_storage_for_context(root: str, context_id: str) -> FileCheckpointStorage: + """Build a ``FileCheckpointStorage`` for a hosted response/conversation context. + + ``context_id`` originates from caller-controlled fields such as + ``previous_response_id`` or from server-generated fields such as + ``conversation_id`` / ``response_id``. In every case it must be treated as + an untrusted single path segment: path separators, drive letters, parent + references and similar would otherwise let the resulting directory escape + the configured checkpoint root (CWE-22). + """ + if not isinstance(context_id, str) or not context_id: + raise RuntimeError("Invalid checkpoint context id: must be a non-empty string.") + # Treat every hosted context id as one untrusted path segment. Do not URL-decode here: + # hosting never decodes these ids before joining them, so encoded traversal markers + # are accepted as literal directory names. + if ( + "/" in context_id + or "\\" in context_id + or "\x00" in context_id + or context_id.strip(".") == "" + or os.path.isabs(context_id) + or os.path.splitdrive(context_id)[0] + ): + raise RuntimeError(f"Invalid checkpoint context id: {context_id!r}") + + root_path = Path(root).resolve() + storage_path = (root_path / context_id).resolve() + if not storage_path.is_relative_to(root_path): + raise RuntimeError(f"Invalid checkpoint context id: {context_id!r}") + return FileCheckpointStorage( + storage_path, + # Hosted workflow checkpoints can persist Azure's role enum inside Message objects. + allowed_checkpoint_types=[AZURE_RESPONSES_MESSAGE_ROLE_TYPE], + ) + + +_AZURE_RESPONSES_MESSAGE_ROLE_TYPE = AZURE_RESPONSES_MESSAGE_ROLE_TYPE +_checkpoint_storage_for_context = checkpoint_storage_for_context + + +@dataclass(frozen=True) +class HostedWorkflowConversationTurn: + """Workflow agent instance and checkpoint resources for one hosted response turn.""" + + agent: WorkflowAgent + restore_context_id: str | None + write_context_id: str + restore_checkpoint_id: str | None + restore_checkpoint_storage: FileCheckpointStorage | None + write_checkpoint_storage: FileCheckpointStorage + + async def run_non_streaming(self, input_messages: Sequence[Message]) -> AgentResponse: + """Run this hosted workflow turn in non-streaming mode.""" + await self._restore_checkpoint(stream=False) + return await self.agent.run( + input_messages, + stream=False, + checkpoint_storage=self.write_checkpoint_storage, + ) + + async def run_streaming(self, input_messages: Sequence[Message]) -> AsyncIterator[AgentResponseUpdate]: + """Run this hosted workflow turn in streaming mode.""" + await self._restore_checkpoint(stream=True) + async for update in self.agent.run( + input_messages, + stream=True, + checkpoint_storage=self.write_checkpoint_storage, + ): + yield update + + async def _restore_checkpoint(self, *, stream: bool) -> None: + """Restore the previous workflow checkpoint, if this turn has one. + + Hosted ``previous_response_id`` turns restore from the previous response's + checkpoint context but write new checkpoints under the current response. + ``WorkflowAgent.run`` accepts one checkpoint storage, so restoration and + input execution remain separate core calls inside this adapter. + """ + if self.restore_checkpoint_id is None: + return + if self.restore_checkpoint_storage is None: # pragma: no cover - defensive invariant + raise RuntimeError("Restore checkpoint storage is not configured.") + + if stream: + async for _ in self.agent.run( + stream=True, + checkpoint_id=self.restore_checkpoint_id, + checkpoint_storage=self.restore_checkpoint_storage, + ): + pass + return + + await self.agent.run( + stream=False, + checkpoint_id=self.restore_checkpoint_id, + checkpoint_storage=self.restore_checkpoint_storage, + ) + + async def delete_not_latest_checkpoints(self, workflow_name: str) -> None: + """Delete old checkpoints from this turn's write context.""" + latest_checkpoint = await self.write_checkpoint_storage.get_latest(workflow_name=workflow_name) + if latest_checkpoint is None: + return + all_checkpoints = await self.write_checkpoint_storage.list_checkpoints(workflow_name=workflow_name) + for checkpoint in all_checkpoints: + if checkpoint.checkpoint_id != latest_checkpoint.checkpoint_id: + await self.write_checkpoint_storage.delete(checkpoint.checkpoint_id) + + +def copy_workflow_agent_for_hosted_turn(agent: WorkflowAgent) -> WorkflowAgent: + """Create a fresh workflow agent instance for a hosted workflow turn.""" + try: + return copy.deepcopy(agent) + except Exception as exc: + raise RuntimeError( + "Hosted workflow agents must be copyable so each response turn can run with isolated workflow state." + ) from exc + + +class HostedWorkflowConversationAdapter: + """Resolves checkpoint contexts for workflow agents hosted behind Responses.""" + + def __init__( + self, + checkpoint_storage_root: str, + workflow_agent_factory: Callable[[], WorkflowAgent], + ) -> None: + self._checkpoint_storage_root = checkpoint_storage_root + self._workflow_agent_factory = workflow_agent_factory + + async def prepare_turn( + self, + *, + response_id: str, + previous_response_id: str | None, + conversation_id: str | None, + ) -> HostedWorkflowConversationTurn: + """Prepare restore and write checkpoint storage for one hosted workflow turn. + + ``previous_response_id`` restores from the prior response and writes + under the current response. ``conversation_id`` restores and writes + under the same stable conversation context. + """ + if previous_response_id is not None and conversation_id is not None: + raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.") + + restore_context_id = previous_response_id or conversation_id + restore_checkpoint_id: str | None = None + restore_checkpoint_storage: FileCheckpointStorage | None = None + if restore_context_id is not None: + restore_checkpoint_storage = checkpoint_storage_for_context( + self._checkpoint_storage_root, + restore_context_id, + ) + + write_context_id = conversation_id or response_id + write_checkpoint_storage = checkpoint_storage_for_context(self._checkpoint_storage_root, write_context_id) + + agent = self._workflow_agent_factory() + if not isinstance(agent, WorkflowAgent): + raise RuntimeError("Workflow agent factory did not return a WorkflowAgent.") + + if restore_checkpoint_storage is not None: + latest_checkpoint = await restore_checkpoint_storage.get_latest(workflow_name=agent.workflow.name) + if latest_checkpoint is not None: + restore_checkpoint_id = latest_checkpoint.checkpoint_id + + return HostedWorkflowConversationTurn( + agent=agent, + restore_context_id=restore_context_id, + write_context_id=write_context_id, + restore_checkpoint_id=restore_checkpoint_id, + restore_checkpoint_storage=restore_checkpoint_storage, + write_checkpoint_storage=write_checkpoint_storage, + ) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 8165b78342..959979c4be 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -12,14 +12,12 @@ from collections.abc import AsyncIterable, AsyncIterator, Generator, Mapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack, suppress from dataclasses import asdict, dataclass, is_dataclass -from pathlib import Path from typing import Protocol, cast from agent_framework import ( ChatOptions, Content, ContextProvider, - FileCheckpointStorage, HistoryProvider, Message, RawAgent, @@ -72,7 +70,6 @@ MessageContentOutputTextContent, MessageContentReasoningTextContent, MessageContentRefusalContent, - MessageRole, OAuthConsentRequestOutputItem, OutputItem, OutputItemApplyPatchToolCall, @@ -115,9 +112,17 @@ from mcp import McpError from typing_extensions import Any +from ._hosted_workflow_conversation import ( + AZURE_RESPONSES_MESSAGE_ROLE_TYPE, + HostedWorkflowConversationAdapter, + checkpoint_storage_for_context, + copy_workflow_agent_for_hosted_turn, +) + logger = logging.getLogger(__name__) -_AZURE_RESPONSES_MESSAGE_ROLE_TYPE = f"{MessageRole.__module__}:{MessageRole.__qualname__}" +_AZURE_RESPONSES_MESSAGE_ROLE_TYPE = AZURE_RESPONSES_MESSAGE_ROLE_TYPE +_checkpoint_storage_for_context = checkpoint_storage_for_context # region Approval Storage @@ -214,52 +219,6 @@ async def load_approval_request(self, approval_request_id: str) -> Content: return await asyncio.to_thread(self._load_sync, approval_request_id) -def _checkpoint_storage_for_context(root: str, context_id: str) -> FileCheckpointStorage: - """Build a ``FileCheckpointStorage`` for ``context_id`` rooted under ``root``. - - ``context_id`` originates from caller-controlled fields such as - ``previous_response_id`` or from server-generated fields such as - ``conversation_id`` / ``response_id``. In every case it must be treated as - an untrusted single path segment: path separators, drive letters, parent - references and similar would otherwise let the resulting directory escape - the configured checkpoint root (CWE-22). The check resolves the joined - path and verifies it stays under the resolved root before any directory is - created on disk. - """ - if not isinstance(context_id, str) or not context_id: - raise RuntimeError("Invalid checkpoint context id: must be a non-empty string.") - # Reject any segment that is not a single safe path component. This covers - # POSIX/Windows separators, NUL bytes, drive letters, and all-dot segments - # (``.``, ``..``, ``...``, ...). We deliberately do not URL-decode the id - # here: the hosting layer never decodes context ids before joining them, so - # forms such as ``%2e%2e`` are accepted as literal directory names. Do NOT - # add decoding here without re-validating after the decode -- decode-then- - # join is exactly the pattern that reintroduces traversal. We also do not - # attempt to "sanitize" by stripping characters because that can introduce - # collisions between distinct ids. - if ( - "/" in context_id - or "\\" in context_id - or "\x00" in context_id - # All-dot segments (``.``, ``..``, ``...``, ...) reduce to "" after stripping dots. - or context_id.strip(".") == "" - or os.path.isabs(context_id) - or os.path.splitdrive(context_id)[0] - ): - raise RuntimeError(f"Invalid checkpoint context id: {context_id!r}") - - root_path = Path(root).resolve() - storage_path = (root_path / context_id).resolve() - if not storage_path.is_relative_to(root_path): - raise RuntimeError(f"Invalid checkpoint context id: {context_id!r}") - return FileCheckpointStorage( - storage_path, - # Keep this provider-specific allowlist narrow. Hosted workflow - # checkpoints can persist Azure's role enum inside Message objects. - allowed_checkpoint_types=[_AZURE_RESPONSES_MESSAGE_ROLE_TYPE], - ) - - # endregion Approval Storage # Foundry Toolbox Auth integration @@ -574,89 +533,32 @@ async def _handle_inner_workflow( if are_options_set: logger.warning("Workflow agent doesn't support runtime options. They will be ignored.") - if request.previous_response_id is not None and context.conversation_id is not None: - raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.") - context_id = request.previous_response_id or context.conversation_id - # The following should never happen due to the checks above. # This is for type safety and defensive programming. if self._checkpoint_storage_path is None: raise RuntimeError("Checkpoint storage path is not configured for workflow agent.") if not isinstance(self._agent, WorkflowAgent): raise RuntimeError("Agent is not a workflow agent.") + workflow_agent_template = self._agent + workflow_conversation_adapter = HostedWorkflowConversationAdapter( + self._checkpoint_storage_path, + lambda: copy_workflow_agent_for_hosted_turn(workflow_agent_template), + ) # Workflow agents are not async context managers in any built-in path, # but call _ensure_agent_ready for symmetry with the regular path so # any future async resources owned by the workflow are entered here. await self._ensure_agent_ready() - # Determine the latest checkpoint (if any) so we can resume the - # workflow's prior state for this turn. The directory is keyed by - # the inbound context id (conversation_id when set, otherwise - # previous_response_id). Multi-turn declarative workflows need the - # workflow's internal state (e.g. Conversation.messages, - # intermediate Local.* variables) to survive across user turns; - # the only place that state lives is the workflow checkpoint, so - # on every turn we restore the latest checkpoint and feed the new - # input back into the start executor as a continuation rather than - # a fresh run. - latest_checkpoint_id: str | None = None - restore_storage: FileCheckpointStorage | None = None - if context_id is not None: - restore_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, context_id) - latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name) - if latest_checkpoint is not None: - latest_checkpoint_id = latest_checkpoint.checkpoint_id - - # Storage that will receive checkpoints written during this turn. - # When the caller chains with previous_response_id, the next turn - # will reference the current response_id as its previous_response_id, - # so new checkpoints must land under the current response_id (or the - # conversation_id when set). When conversation_id is set, this - # matches restore_storage; when only previous_response_id was - # supplied, restore_storage points at the *prior* response's - # directory and write_storage points at the *current* response's. - write_context_id = context.conversation_id or context.response_id - write_storage = _checkpoint_storage_for_context(self._checkpoint_storage_path, write_context_id) - - # Multi-turn pattern: when we have a prior checkpoint, restore it - # first (drive the workflow back to idle with prior state intact), - # then make a separate call that delivers the new user input. This - # depends on Workflow.run preserving shared state across calls. The - # restore-only call may yield events from any pending in-flight - # work in the checkpoint; we consume those internally here so they - # don't surface to the response stream as duplicates. - # - # If the restored checkpoint had pending request_info events, the - # restore-only call replays them through - # ``WorkflowAgent._convert_workflow_event_to_agent_response_updates`` - # and populates ``self._agent.pending_requests``. That is the correct - # state: those requests are genuinely outstanding, and the next - # ``run(input_messages, ...)`` call may contain ``function_call_output`` - # items (carried as FunctionResult/FunctionApprovalResponse content) - # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`. - if latest_checkpoint_id is not None: - if is_streaming_request: - async for _ in self._agent.run( - stream=True, - checkpoint_id=latest_checkpoint_id, - checkpoint_storage=restore_storage, - ): - pass - else: - await self._agent.run( - stream=False, - checkpoint_id=latest_checkpoint_id, - checkpoint_storage=restore_storage, - ) + turn = await workflow_conversation_adapter.prepare_turn( + response_id=context.response_id, + previous_response_id=request.previous_response_id, + conversation_id=context.conversation_id, + ) if not is_streaming_request: # Run the agent in non-streaming mode with the new user input. - response = await self._agent.run( - input_messages, - stream=False, - checkpoint_storage=write_storage, - ) + response = await turn.run_non_streaming(input_messages) async for item in _to_outputs_for_messages( response_event_stream, @@ -665,18 +567,14 @@ async def _handle_inner_workflow( ): yield item - await self._delete_not_latest_checkpoints(write_storage, self._agent.workflow.name) + await turn.delete_not_latest_checkpoints(turn.agent.workflow.name) yield response_event_stream.emit_completed() return tracker = _OutputItemTracker(response_event_stream) # Run the workflow agent in streaming mode with the new user input. - async for update in self._agent.run( - input_messages, - stream=True, - checkpoint_storage=write_storage, - ): + async for update in turn.run_streaming(input_messages): for content in update.contents: for event in tracker.handle(content): yield event @@ -691,26 +589,13 @@ async def _handle_inner_workflow( for event in tracker.close(): yield event - await self._delete_not_latest_checkpoints(write_storage, self._agent.workflow.name) + await turn.delete_not_latest_checkpoints(turn.agent.workflow.name) yield response_event_stream.emit_completed() except Exception as ex: logger.exception("Failed to produce response for workflow agent") for event in self._emit_failure(response_event_stream, tracker, ex): yield event - @staticmethod - async def _delete_not_latest_checkpoints(checkpoint_storage: FileCheckpointStorage, workflow_name: str) -> None: - """Delete all checkpoints except the latest one. - - We only need the last checkpoint for each invocation. - """ - latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=workflow_name) - if latest_checkpoint is not None: - all_checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow_name) - for checkpoint in all_checkpoints: - if checkpoint.checkpoint_id != latest_checkpoint.checkpoint_id: - await checkpoint_storage.delete(checkpoint.checkpoint_id) - @staticmethod def _emit_failure( response_event_stream: ResponseEventStream, diff --git a/python/packages/foundry_hosting/tests/test_hosted_workflow_conversation.py b/python/packages/foundry_hosting/tests/test_hosted_workflow_conversation.py new file mode 100644 index 0000000000..4f26c6e9d4 --- /dev/null +++ b/python/packages/foundry_hosting/tests/test_hosted_workflow_conversation.py @@ -0,0 +1,179 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for hosted workflow conversation checkpoint planning.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from agent_framework import ( + Content, + Message, + WorkflowAgent, + WorkflowBuilder, + WorkflowCheckpoint, + WorkflowContext, + executor, +) +from typing_extensions import Any + +from agent_framework_foundry_hosting._hosted_workflow_conversation import ( # pyright: ignore[reportPrivateUsage] + HostedWorkflowConversationAdapter, + checkpoint_storage_for_context, +) + + +async def test_fresh_turn_writes_current_response_context(tmp_path: Path) -> None: + adapter = _adapter(tmp_path) + + turn = await adapter.prepare_turn( + response_id="resp-1", + previous_response_id=None, + conversation_id=None, + ) + + assert isinstance(turn.agent, WorkflowAgent) + assert turn.restore_context_id is None + assert turn.restore_checkpoint_id is None + assert turn.restore_checkpoint_storage is None + assert turn.write_context_id == "resp-1" + assert turn.write_checkpoint_storage.storage_path == (tmp_path / "resp-1").resolve() + + +async def test_previous_response_id_turn_restores_previous_context_and_writes_current_response_context( + tmp_path: Path, +) -> None: + await _save_checkpoint(tmp_path, "resp-1", "ckpt-1") + adapter = _adapter(tmp_path) + + turn = await adapter.prepare_turn( + response_id="resp-2", + previous_response_id="resp-1", + conversation_id=None, + ) + + assert turn.restore_context_id == "resp-1" + assert turn.restore_checkpoint_id == "ckpt-1" + assert turn.restore_checkpoint_storage is not None + assert turn.restore_checkpoint_storage.storage_path == (tmp_path / "resp-1").resolve() + assert turn.write_context_id == "resp-2" + assert turn.write_checkpoint_storage.storage_path == (tmp_path / "resp-2").resolve() + + +async def test_conversation_id_turn_restores_and_writes_same_stable_context(tmp_path: Path) -> None: + await _save_checkpoint(tmp_path, "conv-alpha", "ckpt-1") + adapter = _adapter(tmp_path) + + turn = await adapter.prepare_turn( + response_id="resp-2", + previous_response_id=None, + conversation_id="conv-alpha", + ) + + assert turn.restore_context_id == "conv-alpha" + assert turn.restore_checkpoint_id == "ckpt-1" + assert turn.restore_checkpoint_storage is not None + assert turn.restore_checkpoint_storage.storage_path == (tmp_path / "conv-alpha").resolve() + assert turn.write_context_id == "conv-alpha" + assert turn.write_checkpoint_storage.storage_path == (tmp_path / "conv-alpha").resolve() + + +async def test_previous_response_id_and_conversation_id_are_mutually_exclusive(tmp_path: Path) -> None: + adapter = _adapter(tmp_path) + + with pytest.raises(RuntimeError, match="Previous response ID cannot be used in conjunction with conversation ID"): + await adapter.prepare_turn( + response_id="resp-2", + previous_response_id="resp-1", + conversation_id="conv-alpha", + ) + + +async def test_turn_prunes_old_checkpoints_from_write_context(tmp_path: Path) -> None: + storage = checkpoint_storage_for_context(str(tmp_path), "conv-alpha") + await storage.save( + WorkflowCheckpoint( + workflow_name="wf", + graph_signature_hash="hash", + checkpoint_id="ckpt-old", + timestamp="2024-01-01T00:00:00+00:00", + ) + ) + await storage.save( + WorkflowCheckpoint( + workflow_name="wf", + graph_signature_hash="hash", + checkpoint_id="ckpt-new", + timestamp="2024-01-02T00:00:00+00:00", + ) + ) + adapter = _adapter(tmp_path) + turn = await adapter.prepare_turn( + response_id="resp-1", + previous_response_id=None, + conversation_id="conv-alpha", + ) + + await turn.delete_not_latest_checkpoints("wf") + + checkpoint_ids = sorted( + checkpoint.checkpoint_id + for checkpoint in await turn.write_checkpoint_storage.list_checkpoints(workflow_name="wf") + ) + assert checkpoint_ids == ["ckpt-new"] + + +async def test_adapter_uses_new_workflow_agent_for_each_turn(tmp_path: Path) -> None: + created_agents: list[WorkflowAgent] = [] + + def factory() -> WorkflowAgent: + agent = _build_workflow_agent() + created_agents.append(agent) + return agent + + adapter = HostedWorkflowConversationAdapter(str(tmp_path), factory) + + first = await adapter.prepare_turn( + response_id="resp-1", + previous_response_id=None, + conversation_id=None, + ) + second = await adapter.prepare_turn( + response_id="resp-2", + previous_response_id=None, + conversation_id=None, + ) + + assert first.agent is created_agents[0] + assert second.agent is created_agents[1] + assert first.agent is not second.agent + + +def _adapter(root: Path) -> HostedWorkflowConversationAdapter: + return HostedWorkflowConversationAdapter(str(root), _build_workflow_agent) + + +def _build_workflow_agent() -> WorkflowAgent: + @executor + async def start(messages: list[Message], ctx: WorkflowContext[Any, Message]) -> None: + await ctx.yield_output(Message(role="assistant", contents=[Content.from_text("ok")])) + + return WorkflowBuilder(name="wf", start_executor=start, output_from=[start]).build().as_agent() + + +async def _save_checkpoint( + root: Path, + context_id: str, + checkpoint_id: str, + *, + workflow_name: str = "wf", +) -> None: + storage = checkpoint_storage_for_context(str(root), context_id) + await storage.save( + WorkflowCheckpoint( + workflow_name=workflow_name, + graph_signature_hash="hash", + checkpoint_id=checkpoint_id, + ) + ) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index eabf4231a7..2639f7f08d 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -3077,7 +3077,13 @@ async def test_handle_inner_workflow_restores_message_role_checkpoint_from_previ ) input_item = ItemMessage({"type": "message", "role": "user", "content": "next turn"}) - with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])): + with ( + patch( + "agent_framework_foundry_hosting._responses.copy_workflow_agent_for_hosted_turn", + return_value=agent, + ), + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[input_item])), + ): async for _ in server._handle_inner_workflow(request, context): # pyright: ignore[reportPrivateUsage] pass @@ -3689,6 +3695,9 @@ def __init__( self.run_count = 0 self.last_run_messages: list[Message] = [] + def __deepcopy__(self, memo: dict[int, Any]) -> _ToolApprovalWorkflowAgentMock: + return self + def create_session(self, **kwargs: Any) -> AgentSession: return AgentSession() @@ -3884,6 +3893,29 @@ async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorReque return WorkflowAgent(workflow=workflow, name="Text Workflow Agent") +def _build_stateful_workflow_agent() -> WorkflowAgent: + """Build a workflow agent that exposes shared workflow state in its output.""" + + @executor + async def start(messages: list[Message], ctx: WorkflowContext[Any, Message]) -> None: + turn_count = int(ctx.get_state("turn_count", 0)) + 1 + ctx.set_state("turn_count", turn_count) + await ctx.yield_output(Message(role="assistant", contents=[Content.from_text(f"turn_count={turn_count}")])) + + workflow = WorkflowBuilder(name="Stateful Workflow", start_executor=start, output_from=[start]).build() + return WorkflowAgent(workflow=workflow, name="Stateful Workflow Agent") + + +def _output_texts(body: dict[str, Any]) -> list[str]: + return [ + part["text"] + for item in body["output"] + if item["type"] == "message" + for part in item.get("content", []) + if part.get("type") == "output_text" + ] + + def _build_approval_workflow_agent( *, approval_request_id: str, @@ -3949,6 +3981,39 @@ async def test_basic_text_response_streaming(self) -> None: text_done = [e for e in events if e["event"] == "response.output_text.done"] assert any(e["data"]["text"] == "hello stream" for e in text_done) + async def test_fresh_requests_do_not_share_workflow_state(self) -> None: + workflow_agent = _build_stateful_workflow_agent() + server = _make_server(workflow_agent) + + first = await _post(server, input_text="first", stream=False) + second = await _post(server, input_text="second", stream=False) + + assert first.status_code == 200 + assert second.status_code == 200 + assert _output_texts(first.json()) == ["turn_count=1"] + assert _output_texts(second.json()) == ["turn_count=1"] + + async def test_previous_response_id_continues_workflow_state_from_checkpoint(self) -> None: + workflow_agent = _build_stateful_workflow_agent() + server = _make_server(workflow_agent) + + first = await _post(server, input_text="first", stream=False) + first_body = first.json() + second = await _post_json( + server, + { + "model": "test-model", + "input": "second", + "stream": False, + "previous_response_id": first_body["id"], + }, + ) + + assert first.status_code == 200 + assert second.status_code == 200 + assert _output_texts(first_body) == ["turn_count=1"] + assert _output_texts(second.json()) == ["turn_count=2"] + async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None: workflow_agent, mock_agent = _build_approval_workflow_agent(approval_request_id="apr_wf_ns") server = _make_server(workflow_agent)