Python: Allow branching from hosted Foundry conversations

python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py

@@ -14,6 +14,7 @@
 
 from agent_framework import (
     ChatOptions,
+    CheckpointID,
     CheckpointStorage,
     Content,
     ContextProvider,
@@ -24,10 +25,11 @@
     SessionStore,
     SupportsAgentRun,
     WorkflowAgent,
+    WorkflowCheckpoint,
 )
 from agent_framework._telemetry import mark_feature_used
 from agent_framework.exceptions import AgentFrameworkException
-from azure.ai.agentserver.core import get_request_context
+from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context
 from azure.ai.agentserver.responses import (
     ResponseContext,
     ResponseProviderProtocol,
@@ -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 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:
+        self._storage = storage
+        self.latest_checkpoint_id: CheckpointID | None = None
+
+    async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
+        checkpoint_id = await self._storage.save(checkpoint)
+        self.latest_checkpoint_id = checkpoint_id
+        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 (
@@ -334,8 +369,9 @@ async def _handle_inner_agent(
     ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]:
         """Handle a regular agent with Responses-managed MAF session continuity.
 
-        Conversation mode reads and writes one MAF session snapshot under
-        ``conversation_id``. Response chaining reads the snapshot under
+        Conversation mode 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. Hosted storage uses the request user as its isolation boundary.
@@ -387,6 +423,8 @@ async def _handle_inner_agent(
             )
 
             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.")
             session_load_id = context.conversation_id or previous_response_id
             session = await session_storage.get(session_load_id) if session_load_id is not None else None
             if session is None:
@@ -395,7 +433,6 @@ async def _handle_inner_agent(
                         f"Cannot find an existing agent session for previous_response_id={previous_response_id}."
                     )
                 session = self._agent.create_session()
-            session_save_id = context.conversation_id or context.response_id
         except Exception as ex:
             logger.error("Failed to prepare state storage: %s", ex, exc_info=(type(ex), ex, ex.__traceback__))
             for event in self._emit_failure(response_event_stream, None, ex):
@@ -455,17 +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(session_save_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(
@@ -517,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,
             )
 
@@ -542,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,
@@ -554,57 +612,123 @@ async def _handle_inner_workflow(
                         f"Cannot find an existing workflow checkpoint for previous_response_id={previous_response_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 is not None:
-                async for _ in self._agent.run(
+            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
+                # 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 is not None:
+                    async for _ in self._agent.run(
+                        stream=True,
+                        checkpoint_id=latest_checkpoint.checkpoint_id,
+                        checkpoint_storage=restore_checkpoint_storage,
+                    ):
+                        pass
+
+                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_id=latest_checkpoint.checkpoint_id,
-                    checkpoint_storage=restore_checkpoint_storage,
+                    checkpoint_storage=write_checkpoint_storage,
                 ):
-                    pass
-
-            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=checkpoint_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
+                    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
+
+                # Close any remaining active builder
+                for event in tracker.close():
+                    yield event
+            except (asyncio.CancelledError, GeneratorExit):
+                request_interrupted = True
+                raise
+            except Exception as ex:
+                request_failure = ex
+                logger.error(
+                    "Failed to produce response for workflow agent",
+                    exc_info=(type(ex), ex, ex.__traceback__),
+                )
+            finally:
+                try:
+                    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._copy_workflow_checkpoint_to_conversation(
+                            checkpoint,
+                            conversation_id=context.conversation_id,
+                            platform_context=request_context,
+                        )
+                except Exception as save_error:
+                    save_failure = save_error
+                    if request_interrupted:
+                        message = "Failed to persist the workflow checkpoint while unwinding an interrupted request"
+                    elif request_failure is not None:
+                        message = "Failed to persist the workflow checkpoint after a workflow failure"
+                    else:
+                        message = "Failed to persist the workflow checkpoint after a successful request"
+                    logger.error(message, exc_info=(type(save_error), save_error, save_error.__traceback__))
 
-            # Close any remaining active builder
-            for event in tracker.close():
-                yield event
-            yield response_event_stream.emit_completed()
+            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 persistence 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 _copy_workflow_checkpoint_to_conversation(
+        self,
+        checkpoint: WorkflowCheckpoint,
+        *,
+        conversation_id: str,
+        platform_context: FoundryAgentRequestContext,
+    ) -> None:
+        """Copy a response turn's latest workflow checkpoint to its conversation."""
+        conversation_storage = self._checkpoint_storage_provider.get_store(
+            config=self.config,
+            context_id=conversation_id,
+            platform_context=platform_context,
+        )
+        await conversation_storage.save(checkpoint)
+
     @staticmethod
     def _emit_failure(
         response_event_stream: ResponseEventStream,