From 6aae03d6e5d58020f47e409552db0db0b5c776b1 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:27:02 +0000 Subject: [PATCH 1/9] Python: Allow branching from hosted Foundry conversations --- python/packages/foundry_hosting/README.md | 34 ++++--- .../_responses.py | 62 +++++++++--- .../_session_store.py | 6 +- .../foundry_hosting/tests/test_responses.py | 98 +++++++++++++++++++ 4 files changed, 171 insertions(+), 29 deletions(-) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index ea8818d1e80..93c63b7e514 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -6,14 +6,16 @@ This package provides the integration of Agent Framework agents and workflows wi agents in addition to the Responses provider's message history. By default it uses the experimental `FoundrySessionStore` under `/.sessions` when hosted and an in-memory `SessionStore` locally. Hosted snapshots are partitioned by the -Agent Server request context's platform user ID. Snapshot filenames use the -Responses `conversation_id` or `response_id`, depending on the continuation -mode. +Agent Server request context's platform user ID. Snapshot filenames use Responses `response_id` values, with an additional +`conversation_id` snapshot that points to the latest state of each stored +conversation. Foundry's session file API exposes the hosted `$HOME` directory as `/`, so the API path `/.sessions` is stored on disk at `$HOME/.sessions`. -Workflow agents continue to use their existing checkpoint storage layout. +Workflow agents use the same continuation model for their checkpoints: every +turn is stored under its `response_id`, and stored conversations also maintain +a `conversation_id` checkpoint alias for their latest turn. ## Foundry session isolation @@ -33,11 +35,14 @@ A Foundry session controls hosted compute and filesystem lifetime and may host multiple users and Responses conversations. The Foundry session ID is not used as the MAF session identifier. -When `conversation_id` is used, the host reads and writes the same snapshot -under that ID. When `previous_response_id` is used, the host reads that response -snapshot, runs the loaded MAF session, and writes the updated snapshot under the -current response's `response_id`. Multiple responses can therefore branch from -one prior response without overwriting its snapshot. +When `conversation_id` is used, the host reads the latest snapshot under that +ID, then writes the updated state under both the current `response_id` and the +conversation ID. This preserves every turn while keeping conversation +continuation pointed at the latest state. When `previous_response_id` is used, +the host reads that response snapshot, runs the loaded MAF session, and writes +the updated snapshot under the current response's `response_id`. Responses can +therefore branch from any prior turn without overwriting its snapshot, +including turns originally created through a conversation. Foundry does not infer the hosted `agent_session_id` from `previous_response_id`. Callers using response chains must also reuse the @@ -46,15 +51,18 @@ same sandbox and `$HOME/.sessions` filesystem. Conversation objects bind to a stable hosted session automatically. Workflow checkpoints and function approvals preserve the existing Foundry -Hosting layout. Hosted paths insert the validated raw platform user ID: +Hosting roots. Hosted paths insert the validated raw platform user ID: ```text -/.checkpoints/// +/.checkpoints/// +/.checkpoints/// /.function_approvals//approval_requests.json ``` -Local workflow checkpoints use `{cwd}/.checkpoints//`, and local -function approvals remain in memory. +The conversation directory is a latest-state alias. Each response directory +retains the final checkpoint selected for that turn, allowing a later request +to branch from it. Local workflow checkpoints use the same layout without +``, and local function approvals remain in memory. Hosted requests require container protocol `2.0.0`. The v2-only request `call_id` is checked before session, checkpoint, or approval storage is used, 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 49882462cef..e2ecb1461ba 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -569,7 +569,9 @@ async def _handle_inner_agent( Foundry sessions govern hosted compute and filesystem lifetime and may serve multiple users and Responses conversations. Conversation mode - reads and writes one MAF session snapshot under ``conversation_id``. + reads the latest MAF session snapshot under ``conversation_id`` and + writes each turn under both its immutable ``response_id`` and the + conversation ID. Response chaining reads the snapshot under ``previous_response_id`` and writes the updated session under the current ``response_id``, allowing branches without changing the MAF session's own identifier. The request @@ -688,7 +690,9 @@ async def _handle_inner_agent( session.state.pop(_HOSTED_RESPONSES_HISTORY_SOURCE_ID, None) if session is not None and self._session_store is not None: try: - await self._session_store.set(context.conversation_id or context.response_id, session) + await self._session_store.set(context.response_id, session) + if context.conversation_id is not None: + await self._session_store.set(context.conversation_id, session) except Exception as save_error: if request_interrupted: logger.error( @@ -808,15 +812,10 @@ async def _handle_inner_workflow( 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 + # Each turn writes to response-addressed checkpoint storage. + # Conversation continuation is updated from its latest checkpoint + # after the run. + write_context_id = context.response_id write_storage = _checkpoint_storage_for_context( self._checkpoint_storage_path, write_context_id, @@ -869,7 +868,12 @@ async def _handle_inner_workflow( ): yield item - await self._delete_not_latest_checkpoints(write_storage, self._agent.workflow.name) + await self._finalize_workflow_checkpoints( + write_storage, + workflow_name=self._agent.workflow.name, + conversation_id=context.conversation_id, + user_id=user_id, + ) yield response_event_stream.emit_completed() return @@ -895,13 +899,45 @@ 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 self._finalize_workflow_checkpoints( + write_storage, + workflow_name=self._agent.workflow.name, + conversation_id=context.conversation_id, + user_id=user_id, + ) 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 + async def _finalize_workflow_checkpoints( + self, + response_storage: FileCheckpointStorage, + *, + workflow_name: str, + conversation_id: str | None, + user_id: str | None, + ) -> None: + """Keep one response checkpoint and update the conversation's latest-state alias.""" + await self._delete_not_latest_checkpoints(response_storage, workflow_name) + if conversation_id is None: + return + + latest_checkpoint = await response_storage.get_latest(workflow_name=workflow_name) + if latest_checkpoint is None: + return + if self._checkpoint_storage_path is None: + raise RuntimeError("Checkpoint storage path is not configured for workflow agent.") + + conversation_storage = _checkpoint_storage_for_context( + self._checkpoint_storage_path, + conversation_id, + user_id=user_id, + ) + await conversation_storage.save(latest_checkpoint) + await self._delete_not_latest_checkpoints(conversation_storage, workflow_name) + @staticmethod async def _delete_not_latest_checkpoints(checkpoint_storage: FileCheckpointStorage, workflow_name: str) -> None: """Delete all checkpoints except the latest one. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py index 3fb90154eef..a151faae54f 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py @@ -19,9 +19,9 @@ class FoundrySessionStore(FileSessionStore): A Foundry hosted session controls platform compute and filesystem lifetime and may host multiple users and Responses conversations. A MAF :class:`AgentSession` contains framework context state. Snapshots are keyed - by ``conversation_id`` for stored conversations or by Responses - ``response_id`` for response chains; these storage keys are independent of - the MAF session's own identifier. + by every Responses ``response_id``. Stored conversations also update a + snapshot keyed by ``conversation_id`` as an alias for the latest turn. + These storage keys are independent of the MAF session's own identifier. This implementation currently persists through :class:`FileSessionStore`, with each validated platform user ID as a child directory. The diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index cc1d8633043..a6af9c1c29d 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -685,9 +685,46 @@ async def run_with_state(*args: Any, **kwargs: Any) -> AgentResponse: assert seen_session_ids[2] == seen_session_ids[0] assert (tmp_path / "user-1" / "conversation-1.json").is_file() assert (tmp_path / "user-1" / "conversation-2.json").is_file() + assert (tmp_path / "user-1" / "response-1.json").is_file() + assert (tmp_path / "user-1" / "response-2.json").is_file() + assert (tmp_path / "user-1" / "response-3.json").is_file() assert agent.create_session.call_count == 2 assert all(item.kwargs == {} for item in agent.create_session.call_args_list) + async def test_conversation_response_snapshots_support_branching(self) -> None: + seen_counts: list[int] = [] + + async def run_with_state(*args: Any, **kwargs: Any) -> AgentResponse: + session = kwargs["session"] + assert isinstance(session, AgentSession) + count = int(session.state.get("turn_count", 0)) + 1 + session.state["turn_count"] = count + seen_counts.append(count) + return AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text(f"turn {count}")])]) + + agent = _make_agent() + agent.run = AsyncMock(side_effect=run_with_state) + store = SessionStore() + server = _make_server(agent, session_store=store) + + first = await _post(server, input_text="first", conversation_id="conversation-1") + await _post(server, input_text="second", conversation_id="conversation-1") + branch = await _post(server, input_text="branch", previous_response_id=first.json()["id"]) + + assert first.status_code == 200 + assert branch.status_code == 200 + assert seen_counts == [1, 2, 2] + + conversation_snapshot = await store.get("conversation-1") + first_snapshot = await store.get(first.json()["id"]) + branch_snapshot = await store.get(branch.json()["id"]) + assert conversation_snapshot is not None + assert first_snapshot is not None + assert branch_snapshot is not None + assert conversation_snapshot.state["turn_count"] == 2 + assert first_snapshot.state["turn_count"] == 1 + assert branch_snapshot.state["turn_count"] == 2 + async def test_previous_response_chain_restores_session_state(self) -> None: seen_counts: list[int] = [] seen_session_ids: list[str] = [] @@ -4893,6 +4930,67 @@ 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) + @pytest.mark.parametrize("stream", [False, True]) + async def test_conversation_response_checkpoints_support_branching(self, tmp_path: Path, stream: bool) -> None: + @executor + async def count_turns(messages: list[Message], ctx: WorkflowContext[Any, AgentResponse]) -> None: + del messages + turn_count = int(ctx.get_state("turn_count", 0)) + 1 + ctx.set_state("turn_count", turn_count) + await ctx.yield_output( + AgentResponse(messages=[Message("assistant", [Content.from_text(f"turn {turn_count}")])]) + ) + + workflow_agent = WorkflowAgent( + workflow=WorkflowBuilder(start_executor=count_turns).build(), + name="Counting Workflow Agent", + ) + server = _make_server(workflow_agent) + server._checkpoint_storage_path = str(tmp_path) # pyright: ignore[reportPrivateUsage] + + def response_body(response: httpx.Response) -> dict[str, Any]: + if not stream: + return response.json() + return _parse_sse_events(response.text)[-1]["data"]["response"] + + first = await _post( + server, + input_text="first", + conversation_id="conversation-1", + stream=stream, + ) + second = await _post( + server, + input_text="second", + conversation_id="conversation-1", + stream=stream, + ) + first_body = response_body(first) + branch = await _post( + server, + input_text="branch", + previous_response_id=first_body["id"], + stream=stream, + ) + + assert first.status_code == 200 + assert second.status_code == 200 + assert branch.status_code == 200 + assert first_body["status"] == "completed" + assert (tmp_path / first_body["id"]).is_dir() + assert (tmp_path / response_body(second)["id"]).is_dir() + assert (tmp_path / response_body(branch)["id"]).is_dir() + assert (tmp_path / "conversation-1").is_dir() + + branch_text = [ + part["text"] + for item in response_body(branch)["output"] + if item["type"] == "message" + for part in item.get("content", []) + if part["type"] == "output_text" + ] + assert branch_text == ["turn 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) From 02c34afd5d71d3976afb212a1b131fd8cc0d2f68 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:40:39 +0000 Subject: [PATCH 2/9] Fix formatting in README.md for clarity on snapshot filenames --- python/packages/foundry_hosting/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 93c63b7e514..219da08f292 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -6,9 +6,9 @@ This package provides the integration of Agent Framework agents and workflows wi agents in addition to the Responses provider's message history. By default it uses the experimental `FoundrySessionStore` under `/.sessions` when hosted and an in-memory `SessionStore` locally. Hosted snapshots are partitioned by the -Agent Server request context's platform user ID. Snapshot filenames use Responses `response_id` values, with an additional -`conversation_id` snapshot that points to the latest state of each stored -conversation. +Agent Server request context's platform user ID. Snapshot filenames use +Responses `response_id` values, with an additional `conversation_id` snapshot +that points to the latest state of each stored conversation. Foundry's session file API exposes the hosted `$HOME` directory as `/`, so the API path `/.sessions` is stored on disk at `$HOME/.sessions`. From 95eb679bd1d46502055494c1cc3e4ac3b9da41cf Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:02:05 +0000 Subject: [PATCH 3/9] Reject previous response ID when a conversation ID is provided --- .../_responses.py | 2 ++ .../foundry_hosting/tests/test_responses.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+) 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 e2ecb1461ba..59561bf34e4 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -593,6 +593,8 @@ async def _handle_inner_agent( try: approval_storage = self._approval_storage_for_request() + 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.") read_session_id = context.conversation_id or request.previous_response_id if self._session_store is None: if read_session_id is not None: diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index a6af9c1c29d..c6856f0eaeb 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -501,6 +501,22 @@ async def test_continuation_requires_session_store(self, continuation: str) -> N error = getattr(failed_response, "error", None) assert "Session storage is required" in getattr(error, "message", "") + async def test_previous_response_rejected_with_conversation(self) -> None: + agent = _make_agent() + server = _make_server(agent) + response = await _post( + server, + previous_response_id="caresp_aaaaaaaaaaaaaaaa00" + "1" * 32, + conversation_id="conversation-1", + ) + + assert response.status_code == 200 + body = response.json() + assert body["status"] == "failed" + assert body["error"]["message"] == ("Previous response ID cannot be used in conjunction with conversation ID.") + agent.run.assert_not_called() + agent.create_session.assert_not_called() + async def test_previous_response_requires_existing_snapshot(self, tmp_path: Path) -> None: agent = _make_agent() server = _make_server(agent, session_store=FoundrySessionStore(tmp_path)) From e4d58cb351e0bfd61744b9c4047b56310fa0a4dc Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:37:35 +0000 Subject: [PATCH 4/9] Python: Enhance error handling and checkpoint management in ResponsesHostServer --- .../_responses.py | 171 ++++++++++-------- .../foundry_hosting/tests/test_responses.py | 89 +++++++++ 2 files changed, 183 insertions(+), 77 deletions(-) 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 59561bf34e4..e950377d597 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -824,89 +824,106 @@ async def _handle_inner_workflow( user_id=user_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( + request_failure: Exception | None = None + request_interrupted = False + try: + # 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, + ) + + 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_id=latest_checkpoint_id, - checkpoint_storage=restore_storage, + checkpoint_storage=write_storage, ) - 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, - ) - - async for item in _to_outputs_for_messages( - response_event_stream, - response.messages, - approval_storage=approval_storage, - ): - yield item - - await self._finalize_workflow_checkpoints( - write_storage, - workflow_name=self._agent.workflow.name, - conversation_id=context.conversation_id, - user_id=user_id, - ) - yield response_event_stream.emit_completed() - return + async for item in _to_outputs_for_messages( + response_event_stream, + response.messages, + approval_storage=approval_storage, + ): + yield item + else: + tracker = _OutputItemTracker(response_event_stream) - 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, + ): + for content in update.contents: + for event in tracker.handle(content): + yield event + if tracker.needs_async: + async for item in _to_outputs( + response_event_stream, content, approval_storage=approval_storage + ): + yield item + tracker.needs_async = False - # 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, - ): - for content in update.contents: - for event in tracker.handle(content): + # Close any remaining active builder + for event in tracker.close(): yield event - if tracker.needs_async: - async for item in _to_outputs( - response_event_stream, content, approval_storage=approval_storage - ): - yield item - tracker.needs_async = False - - # Close any remaining active builder - for event in tracker.close(): - yield event - - await self._finalize_workflow_checkpoints( - write_storage, - workflow_name=self._agent.workflow.name, - conversation_id=context.conversation_id, - user_id=user_id, - ) + except asyncio.CancelledError: + request_interrupted = True + raise + except GeneratorExit: + request_interrupted = True + raise + except Exception as ex: + request_failure = ex + raise + finally: + try: + await self._finalize_workflow_checkpoints( + write_storage, + workflow_name=self._agent.workflow.name, + conversation_id=context.conversation_id, + user_id=user_id, + ) + except Exception as save_error: + if request_interrupted: + logger.error( + "Failed to finalize workflow checkpoints while unwinding an interrupted request", + exc_info=(type(save_error), save_error, save_error.__traceback__), + ) + elif request_failure is not None: + logger.error( + "Failed to finalize workflow checkpoints after a workflow failure", + exc_info=(type(save_error), save_error, save_error.__traceback__), + ) + else: + raise yield response_event_stream.emit_completed() except Exception as ex: logger.exception("Failed to produce response for workflow agent") diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index c6856f0eaeb..c04f40a5311 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -5007,6 +5007,95 @@ def response_body(response: httpx.Response) -> dict[str, Any]: ] assert branch_text == ["turn 2"] + async def test_failed_conversation_workflow_promotes_latest_response_checkpoint(self, tmp_path: Path) -> None: + workflow_agent = _build_text_workflow_agent("ignored") + checkpoint = WorkflowCheckpoint( + workflow_name=workflow_agent.workflow.name, + graph_signature_hash="hash", + ) + + async def failing_run(*args: Any, **kwargs: Any) -> AgentResponse: + await kwargs["checkpoint_storage"].save(checkpoint) + raise RuntimeError("workflow failed") + + server = _make_server(workflow_agent) + server._checkpoint_storage_path = str(tmp_path) # pyright: ignore[reportPrivateUsage] + request = CreateResponse(model="m", input="hi") + context = ResponseContext( + response_id="response-1", + conversation_id="conversation-1", + mode_flags=MagicMock(), + ) + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(workflow_agent, "run", side_effect=failing_run), + ): + events = [ + event + async for event in server._handle_inner_workflow( # pyright: ignore[reportPrivateUsage] + request, + context, + ) + ] + + conversation_storage = FileCheckpointStorage(tmp_path / "conversation-1") + latest = await conversation_storage.get_latest(workflow_name=workflow_agent.workflow.name) + assert latest is not None + assert latest.checkpoint_id == checkpoint.checkpoint_id + assert getattr(events[-1], "type", None) == "response.failed" + + @pytest.mark.parametrize("interruption", ["cancel", "close"]) + async def test_interrupted_conversation_workflow_promotes_latest_response_checkpoint( + self, + tmp_path: Path, + interruption: str, + ) -> None: + workflow_agent = _build_text_workflow_agent("ignored") + checkpoint = WorkflowCheckpoint( + workflow_name=workflow_agent.workflow.name, + graph_signature_hash="hash", + ) + + async def updates(checkpoint_storage: FileCheckpointStorage) -> AsyncIterator[AgentResponseUpdate]: + await checkpoint_storage.save(checkpoint) + yield AgentResponseUpdate(contents=[Content.from_text("started")], role="assistant") + await asyncio.Event().wait() + + def streaming_run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: + return updates(kwargs["checkpoint_storage"]) + + server = _make_server(workflow_agent) + server._checkpoint_storage_path = str(tmp_path) # pyright: ignore[reportPrivateUsage] + request = CreateResponse(model="m", input="hi", stream=True) + context = ResponseContext( + response_id="response-1", + conversation_id="conversation-1", + mode_flags=MagicMock(), + ) + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(workflow_agent, "run", side_effect=streaming_run), + ): + handler = cast( + AsyncGenerator[Any, None], + server._handle_inner_workflow(request, context), # pyright: ignore[reportPrivateUsage] + ) + await anext(handler) + await anext(handler) + await anext(handler) + if interruption == "cancel": + with pytest.raises(asyncio.CancelledError): + await handler.athrow(asyncio.CancelledError()) + else: + await handler.aclose() + + conversation_storage = FileCheckpointStorage(tmp_path / "conversation-1") + latest = await conversation_storage.get_latest(workflow_name=workflow_agent.workflow.name) + assert latest is not None + assert latest.checkpoint_id == checkpoint.checkpoint_id + 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) From 4703b15190773cb150fd61d7173a285ce7ef7135 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Mon, 10 Aug 2026 09:21:55 +0000 Subject: [PATCH 5/9] Enhance ResponsesHostServer to validate previous_response_id against existing workflow checkpoints --- .../_responses.py | 11 +++++++++-- .../foundry_hosting/tests/test_responses.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) 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 c341908dd84..5549465eb7a 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -511,9 +511,10 @@ async def _handle_inner_workflow( if are_options_set: logger.warning("Workflow agent doesn't support runtime options. They will be ignored.") - if request.get("previous_response_id") is not None and context.conversation_id is not None: + previous_response_id = request.get("previous_response_id") + if 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.get("previous_response_id") or context.conversation_id + context_id = previous_response_id or context.conversation_id if not isinstance(self._agent, WorkflowAgent): raise RuntimeError("Agent is not a workflow agent.") @@ -544,6 +545,12 @@ async def _handle_inner_workflow( 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 + elif previous_response_id is not None: + raise RuntimeError( + f"Cannot find an existing workflow checkpoint for " + f"previous_response_id={previous_response_id}. " + "Ensure that the previous response was created successfully and that the ID is correct." + ) # Storage that will receive checkpoints written during this turn. # Every turn writes under its current response_id so a later request diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 5c002f3333f..22e3ce8f286 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -4361,6 +4361,22 @@ 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_previous_response_requires_existing_workflow_checkpoint(self) -> None: + workflow_agent = _build_text_workflow_agent("should not run") + server = _make_server(workflow_agent) + missing_response_id = "caresp_aaaaaaaaaaaaaaaa00" + "1" * 32 + + response = await _post(server, previous_response_id=missing_response_id) + + assert response.status_code == 200 + body = response.json() + assert body["status"] == "failed" + assert ( + f"Cannot find an existing workflow checkpoint for previous_response_id={missing_response_id}." + in body["error"]["message"] + ) + assert body["output"] == [] + @pytest.mark.parametrize("stream", [False, True]) async def test_conversation_response_checkpoints_support_branching(self, stream: bool) -> None: @executor From 389a0bf690437bc2294bf81edba53d2a87a2be25 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Mon, 17 Aug 2026 10:12:27 +0000 Subject: [PATCH 6/9] refactor: streamline checkpoint handling in ResponsesHostServer and improve error logging --- .../_responses.py | 148 ++++++------ .../foundry_hosting/tests/test_responses.py | 213 +++++++----------- 2 files changed, 158 insertions(+), 203 deletions(-) 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 2c01a525dae..dde965837aa 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -513,11 +513,6 @@ async def _handle_inner_workflow( if are_options_set: logger.warning("Workflow agent doesn't support runtime options. They will be ignored.") - previous_response_id = request.get("previous_response_id") - if 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 = previous_response_id or context.conversation_id - if not isinstance(self._agent, WorkflowAgent): raise RuntimeError("Agent is not a workflow agent.") @@ -526,6 +521,14 @@ async def _handle_inner_workflow( # any future async resources owned by the workflow are entered here. await self._ensure_agent_ready() + checkpoint_save_id = context.conversation_id or context.response_id + _validate_checkpoint_context_id(checkpoint_save_id) + checkpoint_storage = self._checkpoint_storage_provider.get_store( + config=self.config, + context_id=checkpoint_save_id, + platform_context=request_context, + ) + # Determine the latest checkpoint (if any) so we can resume the # workflow's prior state for this turn. The directory is keyed by # the platform derived context_id. Multi-turn declarative workflows @@ -535,38 +538,30 @@ async def _handle_inner_workflow( # 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_checkpoint_storage: CheckpointStorage | None = None - if context_id is not None: - _validate_checkpoint_context_id(context_id) - restore_checkpoint_storage = self._checkpoint_storage_provider.get_store( - config=self.config, - context_id=context_id, - platform_context=request_context, - ) + if request.get("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.") + previous_response_id = request.get("previous_response_id") + checkpoint_load_id = context.conversation_id or previous_response_id + latest_checkpoint = None + restore_checkpoint_storage = checkpoint_storage + if checkpoint_load_id is not None: + _validate_checkpoint_context_id(checkpoint_load_id) + if checkpoint_load_id != checkpoint_save_id: + restore_checkpoint_storage = self._checkpoint_storage_provider.get_store( + config=self.config, + context_id=checkpoint_load_id, + platform_context=request_context, + ) latest_checkpoint = await restore_checkpoint_storage.get_latest(workflow_name=self._agent.workflow.name) - if latest_checkpoint is not None: - latest_checkpoint_id = latest_checkpoint.checkpoint_id - elif previous_response_id is not None: + if latest_checkpoint is None and previous_response_id is not None: raise RuntimeError( f"Cannot find an existing workflow checkpoint for previous_response_id={previous_response_id}." ) - # Storage that will receive checkpoints written during this turn. - # Every turn writes under its current response_id so a later request - # can branch from that exact state via previous_response_id. When a - # conversation_id is set, its store is updated after the run as a - # latest-state alias while the response snapshot remains unchanged. - write_context_id = context.response_id - _validate_checkpoint_context_id(write_context_id) - checkpoint_storage = self._checkpoint_storage_provider.get_store( - config=self.config, - context_id=write_context_id, - platform_context=request_context, - ) - request_failure: Exception | None = None + save_failure: Exception | None = None request_interrupted = False + try: # Multi-turn pattern: when we have a prior checkpoint, restore it # first (drive the workflow back to idle with prior state intact), @@ -584,10 +579,10 @@ async def _handle_inner_workflow( # ``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 latest_checkpoint is not None: async for _ in self._agent.run( stream=True, - checkpoint_id=latest_checkpoint_id, + checkpoint_id=latest_checkpoint.checkpoint_id, checkpoint_storage=restore_checkpoint_storage, ): pass @@ -618,70 +613,71 @@ async def _handle_inner_workflow( raise except Exception as ex: request_failure = ex - raise + logger.error( + "Failed to produce response for workflow agent", + exc_info=(type(ex), ex, ex.__traceback__), + ) finally: try: - await self._finalize_workflow_checkpoints( - checkpoint_storage, - workflow_name=self._agent.workflow.name, - conversation_id=context.conversation_id, - platform_context=request_context, - ) + if context.conversation_id is not None: + await self._snapshot_conversation_workflow_checkpoint( + checkpoint_storage, + workflow_name=self._agent.workflow.name, + response_id=context.response_id, + previous_checkpoint_id=( + latest_checkpoint.checkpoint_id if latest_checkpoint is not None else None + ), + platform_context=request_context, + ) except Exception as save_error: + save_failure = save_error if request_interrupted: - logger.error( - "Failed to finalize workflow checkpoints while unwinding an interrupted request", - exc_info=(type(save_error), save_error, save_error.__traceback__), - ) + message = "Failed to snapshot the workflow checkpoint while unwinding an interrupted request" elif request_failure is not None: - logger.error( - "Failed to finalize workflow checkpoints after a workflow failure", - exc_info=(type(save_error), save_error, save_error.__traceback__), - ) + message = "Failed to snapshot the workflow checkpoint after a workflow failure" else: - raise - yield response_event_stream.emit_completed() + message = "Failed to snapshot the workflow checkpoint after a successful request" + logger.error(message, exc_info=(type(save_error), save_error, save_error.__traceback__)) + + if request_failure is not None and save_failure is not None: + failure = RuntimeError( + f"Workflow request failed: {str(request_failure) or type(request_failure).__name__}; " + f"checkpoint snapshot also failed: {str(save_failure) or type(save_failure).__name__}" + ) + for event in self._emit_failure(response_event_stream, tracker, failure): + yield event + elif request_failure is not None: + for event in self._emit_failure(response_event_stream, tracker, request_failure): + yield event + elif save_failure is not None: + for event in self._emit_failure(response_event_stream, tracker, save_failure): + yield event + else: + 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 - async def _finalize_workflow_checkpoints( + async def _snapshot_conversation_workflow_checkpoint( self, - response_storage: CheckpointStorage, + conversation_storage: CheckpointStorage, *, workflow_name: str, - conversation_id: str | None, + response_id: str, + previous_checkpoint_id: str | None, platform_context: FoundryAgentRequestContext, ) -> None: - """Keep one response checkpoint and update the conversation's latest-state alias.""" - await self._delete_not_latest_checkpoints(response_storage, workflow_name) - if conversation_id is None: - return - - latest_checkpoint = await response_storage.get_latest(workflow_name=workflow_name) - if latest_checkpoint is None: + """Snapshot a conversation turn's latest workflow checkpoint under its response ID.""" + latest_checkpoint = await conversation_storage.get_latest(workflow_name=workflow_name) + if latest_checkpoint is None or latest_checkpoint.checkpoint_id == previous_checkpoint_id: return - conversation_storage = self._checkpoint_storage_provider.get_store( + response_storage = self._checkpoint_storage_provider.get_store( config=self.config, - context_id=conversation_id, + context_id=response_id, platform_context=platform_context, ) - await conversation_storage.save(latest_checkpoint) - await self._delete_not_latest_checkpoints(conversation_storage, workflow_name) - - @staticmethod - async def _delete_not_latest_checkpoints( - checkpoint_storage: CheckpointStorage, - workflow_name: str, - ) -> None: - """Delete all checkpoints except the latest one.""" - 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) + await response_storage.save(latest_checkpoint) @staticmethod def _emit_failure( diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index ae9e863613e..66564c832ba 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -449,46 +449,6 @@ async def test_previous_response_requires_existing_agent_session(self) -> None: class TestAgentSessionPersistence: - async def test_conversations_are_isolated_and_response_snapshots_are_saved(self) -> None: - seen_counts: list[int] = [] - seen_session_ids: list[str] = [] - - def run_with_state(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: - del args - session = kwargs["session"] - assert isinstance(session, AgentSession) - count = int(session.state.get("turn_count", 0)) + 1 - session.state["turn_count"] = count - seen_counts.append(count) - seen_session_ids.append(session.session_id) - - async def updates() -> AsyncIterator[AgentResponseUpdate]: - yield AgentResponseUpdate(contents=[Content.from_text(f"turn {count}")], role="assistant") - - return ResponseStream(updates(), finalizer=AgentResponse.from_updates) - - agent = _make_agent() - agent.run = MagicMock(side_effect=run_with_state) - store = SessionStore() - server = _make_server(agent, session_store=store) - - first = await _post(server, input_text="first", conversation_id="conversation-1") - second = await _post(server, input_text="other", conversation_id="conversation-2") - third = await _post(server, input_text="continue", conversation_id="conversation-1") - - assert seen_counts == [1, 1, 2] - assert seen_session_ids[0] != seen_session_ids[1] - assert seen_session_ids[2] == seen_session_ids[0] - for snapshot_id in ( - first.json()["id"], - second.json()["id"], - third.json()["id"], - "conversation-1", - "conversation-2", - ): - assert await store.get(snapshot_id) is not None - assert agent.create_session.call_count == 2 - async def test_conversation_response_snapshots_support_branching(self) -> None: seen_counts: list[int] = [] @@ -4336,22 +4296,6 @@ 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_previous_response_requires_existing_workflow_checkpoint(self) -> None: - workflow_agent = _build_text_workflow_agent("should not run") - server = _make_server(workflow_agent) - missing_response_id = "caresp_aaaaaaaaaaaaaaaa00" + "1" * 32 - - response = await _post(server, previous_response_id=missing_response_id) - - assert response.status_code == 200 - body = response.json() - assert body["status"] == "failed" - assert ( - f"Cannot find an existing workflow checkpoint for previous_response_id={missing_response_id}." - in body["error"]["message"] - ) - assert body["output"] == [] - @pytest.mark.parametrize("stream", [False, True]) async def test_conversation_response_checkpoints_support_branching(self, stream: bool) -> None: @executor @@ -4399,105 +4343,120 @@ def response_body(response: httpx.Response) -> dict[str, Any]: ] assert branch_text == ["turn 2"] - async def test_failed_conversation_workflow_promotes_latest_response_checkpoint(self) -> None: + @pytest.mark.parametrize( + ("termination", "save_new_checkpoint", "snapshot_failure"), + [ + pytest.param("failure", True, False, id="failure"), + pytest.param("failure", True, True, id="request-and-snapshot-failure"), + pytest.param("success", True, True, id="snapshot-failure"), + pytest.param("cancel", True, False, id="cancel"), + pytest.param("cancel", True, True, id="cancel-and-snapshot-failure"), + pytest.param("close", True, False, id="close"), + pytest.param("close", True, True, id="close-and-snapshot-failure"), + pytest.param("failure", False, False, id="failure-before-checkpoint"), + ], + ) + async def test_incomplete_conversation_workflow_snapshots_only_new_checkpoints( + self, + termination: str, + save_new_checkpoint: bool, + snapshot_failure: bool, + caplog: pytest.LogCaptureFixture, + ) -> None: workflow_agent = _build_text_workflow_agent("ignored") checkpoint = WorkflowCheckpoint( workflow_name=workflow_agent.workflow.name, graph_signature_hash="hash", ) - - async def updates(checkpoint_storage: CheckpointStorage) -> AsyncIterator[AgentResponseUpdate]: - await checkpoint_storage.save(checkpoint) - raise RuntimeError("workflow failed") - yield # pragma: no cover - - def failing_run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: - del args - return updates(kwargs["checkpoint_storage"]) - server = _make_server(workflow_agent) - request = CreateResponse(model="m", input="hi") - context = ResponseContext( - response_id="response-1", - conversation_id="conversation-1", - mode_flags=MagicMock(), - ) - with ( - patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), - patch.object(workflow_agent, "run", side_effect=failing_run), - ): - events = [ - event - async for event in server._handle_inner_workflow( # pyright: ignore[reportPrivateUsage] - request, - context, + if not save_new_checkpoint: + conversation_storage = server._checkpoint_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] + config=server.config, + context_id="conversation-1", + platform_context=get_request_context(), + ) + await conversation_storage.save( + WorkflowCheckpoint( + workflow_name=workflow_agent.workflow.name, + graph_signature_hash="hash", ) - ] - - conversation_storage = server._checkpoint_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] - config=server.config, - context_id="conversation-1", - platform_context=get_request_context(), - ) - latest = await conversation_storage.get_latest(workflow_name=workflow_agent.workflow.name) - assert latest is not None - assert latest.checkpoint_id == checkpoint.checkpoint_id - assert events[-1].get("type") == "response.failed" - - @pytest.mark.parametrize("interruption", ["cancel", "close"]) - async def test_interrupted_conversation_workflow_promotes_latest_response_checkpoint( - self, - interruption: str, - ) -> None: - workflow_agent = _build_text_workflow_agent("ignored") - checkpoint = WorkflowCheckpoint( - workflow_name=workflow_agent.workflow.name, - graph_signature_hash="hash", - ) + ) async def updates(checkpoint_storage: CheckpointStorage) -> AsyncIterator[AgentResponseUpdate]: - await checkpoint_storage.save(checkpoint) + if save_new_checkpoint: + await checkpoint_storage.save(checkpoint) + if termination == "failure": + raise RuntimeError("workflow failed") yield AgentResponseUpdate(contents=[Content.from_text("started")], role="assistant") - await asyncio.Event().wait() + if termination in {"cancel", "close"}: + await asyncio.Event().wait() - def streaming_run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: + def run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: del args return updates(kwargs["checkpoint_storage"]) - server = _make_server(workflow_agent) request = CreateResponse(model="m", input="hi", stream=True) context = ResponseContext( response_id="response-1", conversation_id="conversation-1", mode_flags=MagicMock(), ) + snapshot = ( + AsyncMock(side_effect=RuntimeError("snapshot failed")) + if snapshot_failure + else AsyncMock(wraps=server._snapshot_conversation_workflow_checkpoint) # pyright: ignore[reportPrivateUsage] + ) with ( patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), - patch.object(workflow_agent, "run", side_effect=streaming_run), + patch.object(workflow_agent, "run", side_effect=run), + patch.object(server, "_snapshot_conversation_workflow_checkpoint", new=snapshot), ): - handler = cast( - AsyncGenerator[Any, None], - server._handle_inner_workflow(request, context), # pyright: ignore[reportPrivateUsage] - ) - await anext(handler) - await anext(handler) - await anext(handler) - if interruption == "cancel": - with pytest.raises(asyncio.CancelledError): - await handler.athrow(asyncio.CancelledError()) + if termination in {"failure", "success"}: + events = [ + event + async for event in server._handle_inner_workflow( # pyright: ignore[reportPrivateUsage] + request, + context, + ) + ] + assert events[-1].get("type") == "response.failed" + failed_event = cast(Mapping[str, Any], events[-1]) + response = cast(Mapping[str, Any], failed_event["response"]) + error = cast(Mapping[str, Any], response["error"]) + error_message = str(error["message"]) + if termination == "failure": + assert "workflow failed" in error_message + if snapshot_failure: + assert "snapshot failed" in error_message else: - await handler.aclose() - - conversation_storage = server._checkpoint_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] + handler = cast( + AsyncGenerator[Any, None], + server._handle_inner_workflow(request, context), # pyright: ignore[reportPrivateUsage] + ) + await anext(handler) + await anext(handler) + await anext(handler) + if termination == "close": + await handler.aclose() + else: + with pytest.raises(asyncio.CancelledError): + await handler.athrow(asyncio.CancelledError()) + + response_storage = server._checkpoint_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] config=server.config, - context_id="conversation-1", + context_id="response-1", platform_context=get_request_context(), ) - latest = await conversation_storage.get_latest(workflow_name=workflow_agent.workflow.name) - assert latest is not None - assert latest.checkpoint_id == checkpoint.checkpoint_id + latest = await response_storage.get_latest(workflow_name=workflow_agent.workflow.name) + if save_new_checkpoint and not snapshot_failure: + assert latest is not None + assert latest.checkpoint_id == checkpoint.checkpoint_id + else: + assert latest is None + if termination in {"cancel", "close"} and snapshot_failure: + assert "while unwinding an interrupted request" in caplog.text 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") @@ -4598,9 +4557,9 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None second_body = second.json() assert second_body["status"] == "completed" assert [call.kwargs["context_id"] for call in get_store.call_args_list] == [ - first_response_id, first_response_id, second_body["id"], + first_response_id, ] # The inner agent must have been resumed (restore replay + new turn). From 7b551c2a67785bcb236bee9f20ef04c1d949414d Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Mon, 17 Aug 2026 11:00:21 +0000 Subject: [PATCH 7/9] feat: implement _CapturingCheckpointStorage for improved checkpoint handling in ResponsesHostServer --- .../_responses.py | 55 ++++++++++++++----- 1 file changed, 41 insertions(+), 14 deletions(-) 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 dde965837aa..4df56e46231 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -14,6 +14,7 @@ from agent_framework import ( ChatOptions, + CheckpointID, CheckpointStorage, Content, ContextProvider, @@ -24,6 +25,7 @@ SessionStore, SupportsAgentRun, WorkflowAgent, + WorkflowCheckpoint, ) from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentFrameworkException @@ -77,6 +79,39 @@ _HOSTED_RESPONSES_HISTORY_SOURCE_ID = "_foundry_responses_history" +class _CapturingCheckpointStorage: + """Delegate storage while retaining the latest successful checkpoint save. + + The workflow runner does not expose the checkpoint it creates. Capturing it + here lets the host snapshot the current turn under its response ID without + rescanning the conversation's full checkpoint history. + """ + + def __init__(self, storage: CheckpointStorage) -> None: + self._storage = storage + self.latest_checkpoint: WorkflowCheckpoint | None = None + + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + checkpoint_id = await self._storage.save(checkpoint) + self.latest_checkpoint = checkpoint + return checkpoint_id + + async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: + return await self._storage.load(checkpoint_id) + + async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: + return await self._storage.list_checkpoints(workflow_name=workflow_name) + + async def delete(self, checkpoint_id: CheckpointID) -> bool: + return await self._storage.delete(checkpoint_id) + + async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: + return await self._storage.get_latest(workflow_name=workflow_name) + + async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: + return await self._storage.list_checkpoint_ids(workflow_name=workflow_name) + + def _validate_checkpoint_context_id(context_id: str) -> None: """Validate that a checkpoint context ID is a single safe path component in case file-based storage is used.""" if ( @@ -561,6 +596,7 @@ async def _handle_inner_workflow( request_failure: Exception | None = None save_failure: Exception | None = None request_interrupted = False + write_checkpoint_storage = _CapturingCheckpointStorage(checkpoint_storage) try: # Multi-turn pattern: when we have a prior checkpoint, restore it @@ -593,7 +629,7 @@ async def _handle_inner_workflow( async for update in self._agent.run( input_messages, stream=True, - checkpoint_storage=checkpoint_storage, + checkpoint_storage=write_checkpoint_storage, ): for content in update.contents: for event in tracker.handle(content): @@ -619,14 +655,10 @@ async def _handle_inner_workflow( ) finally: try: - if context.conversation_id is not None: + if write_checkpoint_storage.latest_checkpoint is not None and context.conversation_id is not None: await self._snapshot_conversation_workflow_checkpoint( - checkpoint_storage, - workflow_name=self._agent.workflow.name, + write_checkpoint_storage.latest_checkpoint, response_id=context.response_id, - previous_checkpoint_id=( - latest_checkpoint.checkpoint_id if latest_checkpoint is not None else None - ), platform_context=request_context, ) except Exception as save_error: @@ -661,23 +693,18 @@ async def _handle_inner_workflow( async def _snapshot_conversation_workflow_checkpoint( self, - conversation_storage: CheckpointStorage, + checkpoint: WorkflowCheckpoint, *, - workflow_name: str, response_id: str, - previous_checkpoint_id: str | None, platform_context: FoundryAgentRequestContext, ) -> None: """Snapshot a conversation turn's latest workflow checkpoint under its response ID.""" - latest_checkpoint = await conversation_storage.get_latest(workflow_name=workflow_name) - if latest_checkpoint is None or latest_checkpoint.checkpoint_id == previous_checkpoint_id: - return response_storage = self._checkpoint_storage_provider.get_store( config=self.config, context_id=response_id, platform_context=platform_context, ) - await response_storage.save(latest_checkpoint) + await response_storage.save(checkpoint) @staticmethod def _emit_failure( From 22182a170c0a3d2abd1310edaf193fdfb13803aa Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Wed, 19 Aug 2026 12:46:25 +0000 Subject: [PATCH 8/9] Fix hosted response snapshot persistence Persist conversation aliases independently from immutable response snapshots, and reload the exact saved workflow checkpoint before branching. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e30d54ef-bdf4-42db-934f-91ab9f067927 --- .../_responses.py | 56 ++++++++++++------ .../foundry_hosting/tests/test_responses.py | 57 +++++++++++++++++++ 2 files changed, 96 insertions(+), 17 deletions(-) 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 895e9da74ce..ad850454c5f 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -89,11 +89,11 @@ class _CapturingCheckpointStorage: def __init__(self, storage: CheckpointStorage) -> None: self._storage = storage - self.latest_checkpoint: WorkflowCheckpoint | None = None + self.latest_checkpoint_id: CheckpointID | None = None async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: checkpoint_id = await self._storage.save(checkpoint) - self.latest_checkpoint = checkpoint + self.latest_checkpoint_id = checkpoint_id return checkpoint_id async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: @@ -492,19 +492,37 @@ async def _handle_inner_agent( finally: if self._uses_hosted_responses_history: session.state.pop(_HOSTED_RESPONSES_HISTORY_SOURCE_ID, None) - try: - await session_storage.set(context.response_id, session) - if context.conversation_id is not None: - await session_storage.set(context.conversation_id, session) - except Exception as save_error: - save_failure = save_error - if request_interrupted: - message = "Failed to persist the Agent Framework session while unwinding an interrupted request" - elif request_failure is not None: - message = "Failed to persist the Agent Framework session after an agent failure" - else: - message = "Failed to persist the Agent Framework session after a successful request" - logger.error(message, exc_info=(type(save_error), save_error, save_error.__traceback__)) + if request_interrupted: + message = "Failed to persist the Agent Framework session while unwinding an interrupted request" + elif request_failure is not None: + message = "Failed to persist the Agent Framework session after an agent failure" + else: + message = "Failed to persist the Agent Framework session after a successful request" + + save_errors: list[tuple[str, Exception]] = [] + session_save_ids = [("response snapshot", context.response_id)] + if context.conversation_id is not None: + session_save_ids.append(("conversation", context.conversation_id)) + for save_target, session_save_id in session_save_ids: + try: + await session_storage.set(session_save_id, session) + except Exception as save_error: + save_errors.append((save_target, save_error)) + logger.error( + "%s (%s)", + message, + save_target, + exc_info=(type(save_error), save_error, save_error.__traceback__), + ) + + if len(save_errors) == 1: + save_failure = save_errors[0][1] + elif save_errors: + details = "; ".join( + f"{save_target}: {str(save_error) or type(save_error).__name__}" + for save_target, save_error in save_errors + ) + save_failure = RuntimeError(f"Multiple session persistence operations failed: {details}") if request_failure is not None and save_failure is not None: failure = RuntimeError( @@ -655,9 +673,13 @@ async def _handle_inner_workflow( ) finally: try: - if write_checkpoint_storage.latest_checkpoint is not None and context.conversation_id is not None: + if ( + write_checkpoint_storage.latest_checkpoint_id is not None + and context.conversation_id is not None + ): + checkpoint = await write_checkpoint_storage.load(write_checkpoint_storage.latest_checkpoint_id) await self._snapshot_conversation_workflow_checkpoint( - write_checkpoint_storage.latest_checkpoint, + checkpoint, response_id=context.response_id, platform_context=request_context, ) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 4f36e0cccc6..8a13c0ce5b2 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -263,6 +263,22 @@ async def set(self, session_id: str, session: AgentSession) -> None: raise OSError("session storage is full") +class _FailingResponseSnapshotStore(SessionStore): + def __init__(self, *, fail_conversation: bool = False) -> None: + super().__init__() + self.fail_conversation = fail_conversation + self.set_attempts: list[str] = [] + + async def set(self, session_id: str, session: AgentSession) -> None: + self.set_attempts.append(session_id) + if session_id == "conversation-1": + if self.fail_conversation: + raise OSError("conversation storage is full") + await super().set(session_id, session) + return + raise OSError("response snapshot storage is full") + + _SESSION_STORE_UNSET = object() @@ -669,6 +685,44 @@ async def updates() -> AsyncIterator[AgentResponseUpdate]: assert "response.completed" not in event_types assert store.set_attempts == 1 + async def test_response_snapshot_failure_still_persists_conversation(self) -> None: + store = _FailingResponseSnapshotStore() + agent = _make_agent() + + def run(*_args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: + session = kwargs["session"] + assert isinstance(session, AgentSession) + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + session.state["run_complete"] = True + yield AgentResponseUpdate(contents=[Content.from_text("done")], role="assistant") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent.run = MagicMock(side_effect=run) + server = _make_server(agent, session_store=store) + + response = await _post(server, conversation_id="conversation-1") + conversation = await store.get("conversation-1") + + assert response.json()["status"] == "failed" + assert response.json()["id"] == store.set_attempts[0] + assert store.set_attempts == [response.json()["id"], "conversation-1"] + assert conversation is not None + assert conversation.state["run_complete"] is True + + async def test_response_and_conversation_save_failures_are_both_reported(self) -> None: + store = _FailingResponseSnapshotStore(fail_conversation=True) + server = _make_server(_make_agent(), session_store=store) + + response = await _post(server, conversation_id="conversation-1") + error_message = response.json()["error"]["message"] + + assert response.json()["status"] == "failed" + assert "response snapshot: response snapshot storage is full" in error_message + assert "conversation: conversation storage is full" in error_message + assert store.set_attempts == [response.json()["id"], "conversation-1"] + async def test_run_and_save_failure_emit_one_combined_failure( self, caplog: pytest.LogCaptureFixture, @@ -4378,6 +4432,7 @@ async def test_incomplete_conversation_workflow_snapshots_only_new_checkpoints( checkpoint = WorkflowCheckpoint( workflow_name=workflow_agent.workflow.name, graph_signature_hash="hash", + state={"nested": {"value": "saved"}}, ) server = _make_server(workflow_agent) @@ -4397,6 +4452,7 @@ async def test_incomplete_conversation_workflow_snapshots_only_new_checkpoints( async def updates(checkpoint_storage: CheckpointStorage) -> AsyncIterator[AgentResponseUpdate]: if save_new_checkpoint: await checkpoint_storage.save(checkpoint) + checkpoint.state["nested"]["value"] = "mutated" if termination == "failure": raise RuntimeError("workflow failed") yield AgentResponseUpdate(contents=[Content.from_text("started")], role="assistant") @@ -4464,6 +4520,7 @@ def run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: if save_new_checkpoint and not snapshot_failure: assert latest is not None assert latest.checkpoint_id == checkpoint.checkpoint_id + assert latest.state["nested"]["value"] == "saved" else: assert latest is None if termination in {"cancel", "close"} and snapshot_failure: From 518bfd3f33930da3749cac017f0287856f1e0e12 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:40:29 +0000 Subject: [PATCH 9/9] Align workflow response and conversation state Persist workflow checkpoints under the immutable response ID first, then copy the exact saved checkpoint to the conversation ID. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e30d54ef-bdf4-42db-934f-91ab9f067927 --- .../_responses.py | 39 ++++++++++--------- .../foundry_hosting/tests/test_responses.py | 39 ++++++++++++------- 2 files changed, 44 insertions(+), 34 deletions(-) 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 ad850454c5f..6405baef943 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -82,9 +82,9 @@ class _CapturingCheckpointStorage: """Delegate storage while retaining the latest successful checkpoint save. - The workflow runner does not expose the checkpoint it creates. Capturing it - here lets the host snapshot the current turn under its response ID without - rescanning the conversation's full checkpoint history. + The workflow runner does not expose the checkpoint it creates. Capturing its + ID lets the host copy the exact persisted response state to the conversation + without rescanning the full checkpoint history. """ def __init__(self, storage: CheckpointStorage) -> None: @@ -574,11 +574,12 @@ async def _handle_inner_workflow( # any future async resources owned by the workflow are entered here. await self._ensure_agent_ready() - checkpoint_save_id = context.conversation_id or context.response_id - _validate_checkpoint_context_id(checkpoint_save_id) + # Persist each turn under its immutable response ID first. For conversation turns, + # the latest successfully saved checkpoint is copied to the conversation ID below. + _validate_checkpoint_context_id(context.response_id) checkpoint_storage = self._checkpoint_storage_provider.get_store( config=self.config, - context_id=checkpoint_save_id, + context_id=context.response_id, platform_context=request_context, ) @@ -599,7 +600,7 @@ async def _handle_inner_workflow( restore_checkpoint_storage = checkpoint_storage if checkpoint_load_id is not None: _validate_checkpoint_context_id(checkpoint_load_id) - if checkpoint_load_id != checkpoint_save_id: + if checkpoint_load_id != context.response_id: restore_checkpoint_storage = self._checkpoint_storage_provider.get_store( config=self.config, context_id=checkpoint_load_id, @@ -678,25 +679,25 @@ async def _handle_inner_workflow( and context.conversation_id is not None ): checkpoint = await write_checkpoint_storage.load(write_checkpoint_storage.latest_checkpoint_id) - await self._snapshot_conversation_workflow_checkpoint( + await self._copy_workflow_checkpoint_to_conversation( checkpoint, - response_id=context.response_id, + conversation_id=context.conversation_id, platform_context=request_context, ) except Exception as save_error: save_failure = save_error if request_interrupted: - message = "Failed to snapshot the workflow checkpoint while unwinding an interrupted request" + message = "Failed to persist the workflow checkpoint while unwinding an interrupted request" elif request_failure is not None: - message = "Failed to snapshot the workflow checkpoint after a workflow failure" + message = "Failed to persist the workflow checkpoint after a workflow failure" else: - message = "Failed to snapshot the workflow checkpoint after a successful request" + message = "Failed to persist the workflow checkpoint after a successful request" logger.error(message, exc_info=(type(save_error), save_error, save_error.__traceback__)) if request_failure is not None and save_failure is not None: failure = RuntimeError( f"Workflow request failed: {str(request_failure) or type(request_failure).__name__}; " - f"checkpoint snapshot also failed: {str(save_failure) or type(save_failure).__name__}" + f"checkpoint persistence also failed: {str(save_failure) or type(save_failure).__name__}" ) for event in self._emit_failure(response_event_stream, tracker, failure): yield event @@ -713,20 +714,20 @@ async def _handle_inner_workflow( for event in self._emit_failure(response_event_stream, tracker, ex): yield event - async def _snapshot_conversation_workflow_checkpoint( + async def _copy_workflow_checkpoint_to_conversation( self, checkpoint: WorkflowCheckpoint, *, - response_id: str, + conversation_id: str, platform_context: FoundryAgentRequestContext, ) -> None: - """Snapshot a conversation turn's latest workflow checkpoint under its response ID.""" - response_storage = self._checkpoint_storage_provider.get_store( + """Copy a response turn's latest workflow checkpoint to its conversation.""" + conversation_storage = self._checkpoint_storage_provider.get_store( config=self.config, - context_id=response_id, + context_id=conversation_id, platform_context=platform_context, ) - await response_storage.save(checkpoint) + await conversation_storage.save(checkpoint) @staticmethod def _emit_failure( diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 8a13c0ce5b2..0343dcf04f7 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -4409,15 +4409,15 @@ def response_body(response: httpx.Response) -> dict[str, Any]: assert branch_text == ["turn 2"] @pytest.mark.parametrize( - ("termination", "save_new_checkpoint", "snapshot_failure"), + ("termination", "save_new_checkpoint", "copy_failure"), [ pytest.param("failure", True, False, id="failure"), - pytest.param("failure", True, True, id="request-and-snapshot-failure"), - pytest.param("success", True, True, id="snapshot-failure"), + pytest.param("failure", True, True, id="request-and-copy-failure"), + pytest.param("success", True, True, id="copy-failure"), pytest.param("cancel", True, False, id="cancel"), - pytest.param("cancel", True, True, id="cancel-and-snapshot-failure"), + pytest.param("cancel", True, True, id="cancel-and-copy-failure"), pytest.param("close", True, False, id="close"), - pytest.param("close", True, True, id="close-and-snapshot-failure"), + pytest.param("close", True, True, id="close-and-copy-failure"), pytest.param("failure", False, False, id="failure-before-checkpoint"), ], ) @@ -4425,7 +4425,7 @@ async def test_incomplete_conversation_workflow_snapshots_only_new_checkpoints( self, termination: str, save_new_checkpoint: bool, - snapshot_failure: bool, + copy_failure: bool, caplog: pytest.LogCaptureFixture, ) -> None: workflow_agent = _build_text_workflow_agent("ignored") @@ -4469,16 +4469,16 @@ def run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: conversation_id="conversation-1", mode_flags=MagicMock(), ) - snapshot = ( - AsyncMock(side_effect=RuntimeError("snapshot failed")) - if snapshot_failure - else AsyncMock(wraps=server._snapshot_conversation_workflow_checkpoint) # pyright: ignore[reportPrivateUsage] + copy_checkpoint = ( + AsyncMock(side_effect=RuntimeError("copy failed")) + if copy_failure + else AsyncMock(wraps=server._copy_workflow_checkpoint_to_conversation) # pyright: ignore[reportPrivateUsage] ) with ( patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), patch.object(workflow_agent, "run", side_effect=run), - patch.object(server, "_snapshot_conversation_workflow_checkpoint", new=snapshot), + patch.object(server, "_copy_workflow_checkpoint_to_conversation", new=copy_checkpoint), ): if termination in {"failure", "success"}: events = [ @@ -4495,8 +4495,8 @@ def run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: error_message = str(error["message"]) if termination == "failure": assert "workflow failed" in error_message - if snapshot_failure: - assert "snapshot failed" in error_message + if copy_failure: + assert "copy failed" in error_message else: handler = cast( AsyncGenerator[Any, None], @@ -4517,13 +4517,22 @@ def run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: platform_context=get_request_context(), ) latest = await response_storage.get_latest(workflow_name=workflow_agent.workflow.name) - if save_new_checkpoint and not snapshot_failure: + if save_new_checkpoint: assert latest is not None assert latest.checkpoint_id == checkpoint.checkpoint_id assert latest.state["nested"]["value"] == "saved" + if not copy_failure: + conversation_storage = server._checkpoint_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] + config=server.config, + context_id="conversation-1", + platform_context=get_request_context(), + ) + conversation_latest = await conversation_storage.get_latest(workflow_name=workflow_agent.workflow.name) + assert conversation_latest is not None + assert conversation_latest.to_dict() == latest.to_dict() else: assert latest is None - if termination in {"cancel", "close"} and snapshot_failure: + if termination in {"cancel", "close"} and copy_failure: assert "while unwinding an interrupted request" in caplog.text async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None: