chore: import upstream snapshot with attribution

python/samples/03-workflows/checkpoint/checkpoint_with_human_in_the_loop.py

@@ -0,0 +1,330 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import os
+import sys
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from agent_framework import (
+    Agent,
+    AgentExecutor,
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+    Executor,
+    FileCheckpointStorage,
+    Message,
+    Workflow,
+    WorkflowBuilder,
+    WorkflowContext,
+    handler,
+    response_handler,
+)
+from agent_framework.foundry import FoundryChatClient
+from azure.identity import AzureCliCredential
+from dotenv import load_dotenv
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+# Load environment variables from .env file
+load_dotenv()
+
+"""
+Sample: Checkpoint + human-in-the-loop quickstart.
+
+This getting-started sample keeps the moving pieces to a minimum:
+
+1. A brief is turned into a consistent prompt for an AI copywriter.
+2. The copywriter (an `AgentExecutor`) drafts release notes.
+3. A reviewer gateway sends a request for approval for every draft.
+4. The workflow records checkpoints between each superstep so you can stop the
+   program, restart later, and optionally pre-supply human answers on resume.
+
+Key concepts demonstrated
+-------------------------
+- Minimal executor pipeline with checkpoint persistence.
+- Human-in-the-loop pause/resume with checkpoint restoration.
+
+Typical pause/resume flow
+-------------------------
+1. Run the workflow until a human approval request is emitted.
+2. If the human is offline, exit the program. A checkpoint with
+   ``status=awaiting human response`` now exists.
+3. Later, restart the script, select that checkpoint, and provide the stored
+   human decision when prompted to pre-supply responses.
+   Doing so applies the answer immediately on resume, so the system does **not**
+   re-emit the same ``.
+"""
+
+# Directory used for the sample's temporary checkpoint files. We isolate the
+# demo artefacts so that repeated runs do not collide with other samples and so
+# the clean-up step at the end of the script can simply delete the directory.
+TEMP_DIR = Path(__file__).with_suffix("").parent / "tmp" / "checkpoints_hitl"
+TEMP_DIR.mkdir(parents=True, exist_ok=True)
+
+
+class BriefPreparer(Executor):
+    """Normalises the user brief and sends a single AgentExecutorRequest."""
+
+    # The first executor in the workflow. By keeping it tiny we make it easier
+    # to reason about the state that will later be captured in the checkpoint.
+    # It is responsible for tidying the human-provided brief and kicking off the
+    # agent run with a deterministic prompt structure.
+
+    def __init__(self, id: str, agent_id: str) -> None:
+        super().__init__(id=id)
+        self._agent_id = agent_id
+
+    @handler
+    async def prepare(self, brief: str, ctx: WorkflowContext[AgentExecutorRequest, str]) -> None:
+        # Collapse errant whitespace so the prompt is stable between runs.
+        normalized = " ".join(brief.split()).strip()
+        if not normalized.endswith("."):
+            normalized += "."
+        # Persist the cleaned brief in workflow state so downstream executors and
+        # future checkpoints can recover the original intent.
+        ctx.set_state("brief", normalized)
+        prompt = (
+            "You are drafting product release notes. Summarise the brief below in two sentences. "
+            "Keep it positive and end with a call to action.\n\n"
+            f"BRIEF: {normalized}"
+        )
+        # Hand the prompt to the writer agent. We always route through the
+        # workflow context so the runtime can capture messages for checkpointing.
+        await ctx.send_message(
+            AgentExecutorRequest(messages=[Message("user", contents=[prompt])], should_respond=True),
+            target_id=self._agent_id,
+        )
+
+
+@dataclass
+class HumanApprovalRequest:
+    """Request sent to the human reviewer."""
+
+    # These fields are intentionally simple because they are serialised into
+    # checkpoints. Keeping them primitive types guarantees the new
+    # `pending_requests_from_checkpoint` helper can reconstruct them on resume.
+    prompt: str = ""
+    draft: str = ""
+    iteration: int = 0
+
+
+class ReviewGateway(Executor):
+    """Routes agent drafts to humans and optionally back for revisions."""
+
+    def __init__(self, id: str, writer_id: str) -> None:
+        super().__init__(id=id)
+        self._writer_id = writer_id
+        self._iteration = 0
+
+    @handler
+    async def on_agent_response(self, response: AgentExecutorResponse, ctx: WorkflowContext) -> None:
+        # Capture the agent output so we can surface it to the reviewer and persist iterations.
+        self._iteration += 1
+
+        # Emit a human approval request.
+        await ctx.request_info(
+            request_data=HumanApprovalRequest(
+                prompt="Review the draft. Reply 'approve' or provide edit instructions.",
+                draft=response.agent_response.text,
+                iteration=self._iteration,
+            ),
+            response_type=str,
+        )
+
+    @response_handler
+    async def on_human_feedback(
+        self,
+        original_request: HumanApprovalRequest,
+        feedback: str,
+        ctx: WorkflowContext[AgentExecutorRequest | str, str],
+    ) -> None:
+        # The `original_request` is the request we sent earlier that is now being answered.
+        reply = feedback.strip()
+
+        if len(reply) == 0 or reply.lower() == "approve":
+            # Workflow is completed when the human approves.
+            await ctx.yield_output(original_request.draft)
+            return
+
+        # Any other response loops us back to the writer with fresh guidance.
+        prompt = (
+            "Revise the launch note. Respond with the new copy only.\n\n"
+            f"Previous draft:\n{original_request.draft}\n\n"
+            f"Human guidance: {reply}"
+        )
+        await ctx.send_message(
+            AgentExecutorRequest(messages=[Message("user", contents=[prompt])], should_respond=True),
+            target_id=self._writer_id,
+        )
+
+    @override
+    async def on_checkpoint_save(self) -> dict[str, Any]:
+        # Save the current iteration count in executor state for checkpointing.
+        return {"iteration": self._iteration}
+
+    @override
+    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+        # Restore the iteration count from executor state during checkpoint recovery.
+        self._iteration = state.get("iteration", 0)
+
+
+def create_workflow(checkpoint_storage: FileCheckpointStorage) -> Workflow:
+    """Assemble the workflow graph used by both the initial run and resume."""
+    # Wire the workflow DAG. Edges mirror the numbered steps described in the
+    # module docstring. Because `WorkflowBuilder` is declarative, reading these
+    # edges is often the quickest way to understand execution order.
+    writer_agent = Agent(
+        client=FoundryChatClient(
+            project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+            model=os.environ["FOUNDRY_MODEL"],
+            credential=AzureCliCredential(),
+        ),
+        instructions="Write concise, warm release notes that sound human and helpful.",
+        name="writer",
+    )
+    writer = AgentExecutor(writer_agent)
+    review_gateway = ReviewGateway(id="review_gateway", writer_id="writer")
+    prepare_brief = BriefPreparer(id="prepare_brief", agent_id="writer")
+
+    workflow_builder = (
+        WorkflowBuilder(max_iterations=6, start_executor=prepare_brief, checkpoint_storage=checkpoint_storage)
+        .add_edge(prepare_brief, writer)
+        .add_edge(writer, review_gateway)
+        .add_edge(review_gateway, writer)  # revisions loop
+    )
+
+    return workflow_builder.build()
+
+
+def prompt_for_responses(requests: dict[str, HumanApprovalRequest]) -> dict[str, str]:
+    """Interactive CLI prompt for any live RequestInfo requests."""
+
+    responses: dict[str, str] = {}
+    for request_id, request in requests.items():
+        print("\n=== Human approval needed ===")
+        print(f"request_id: {request_id}")
+        print(f"Iteration: {request.iteration}")
+        print(request.prompt)
+        print("Draft: \n---\n" + request.draft + "\n---")
+        response = input("Type 'approve' or enter revision guidance (or 'exit' to quit): ").strip()
+        if response.lower() == "exit":
+            raise SystemExit("Stopped by user.")
+        responses[request_id] = response
+
+    return responses
+
+
+async def run_interactive_session(
+    workflow: Workflow,
+    initial_message: str | None = None,
+    checkpoint_id: str | None = None,
+) -> str:
+    """Run the workflow until it either finishes or pauses for human input."""
+
+    requests: dict[str, HumanApprovalRequest] = {}
+    responses: dict[str, str] | None = None
+    completed_output: str | None = None
+
+    while True:
+        if responses:
+            event_stream = workflow.run(stream=True, responses=responses)
+            requests.clear()
+            responses = None
+        else:
+            if initial_message:
+                print(f"\nStarting workflow with brief: {initial_message}\n")
+                event_stream = workflow.run(message=initial_message, stream=True)
+            elif checkpoint_id:
+                print("\nStarting workflow from checkpoint...\n")
+                event_stream = workflow.run(checkpoint_id=checkpoint_id, stream=True)
+            else:
+                raise ValueError("Either initial_message or checkpoint_id must be provided")
+
+        async for event in event_stream:
+            if event.type == "status":
+                print(event)
+            if event.type == "output":
+                completed_output = event.data
+            if event.type == "request_info":
+                if isinstance(event.data, HumanApprovalRequest):
+                    requests[event.request_id] = event.data
+                else:
+                    raise ValueError("Unexpected request data type")
+
+        if completed_output:
+            break
+
+        if requests:
+            responses = prompt_for_responses(requests)
+            continue
+
+        raise RuntimeError("Workflow stopped without completing or requesting input")
+
+    return completed_output
+
+
+async def main() -> None:
+    """Entry point used by both the initial run and subsequent resumes."""
+
+    for file in TEMP_DIR.glob("*.json"):
+        # Start each execution with a clean slate so the demonstration is
+        # deterministic even if the directory had stale checkpoints.
+        file.unlink()
+
+    storage = FileCheckpointStorage(storage_path=TEMP_DIR)
+    workflow = create_workflow(checkpoint_storage=storage)
+
+    brief = (
+        "Introduce our limited edition smart coffee grinder. Mention the $249 price, highlight the "
+        "sensor that auto-adjusts the grind, and invite customers to pre-order on the website."
+    )
+
+    print("Running workflow (human approval required)...")
+    result = await run_interactive_session(workflow, initial_message=brief)
+    print(f"Workflow completed with: {result}")
+
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+    if not checkpoints:
+        print("No checkpoints recorded.")
+        return
+
+    sorted_cps = sorted(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp))
+    print("\nAvailable checkpoints:")
+    for idx, cp in enumerate(sorted_cps):
+        print(f"  [{idx}] id={cp.checkpoint_id} iter={cp.iteration_count}")
+
+    # For the pause/resume demo we typically pick the latest checkpoint whose summary
+    # status reads "awaiting human response" - that is the saved state that proves the
+    # workflow can rehydrate, collect the pending answer, and continue after a break.
+    selection = input("\nResume from which checkpoint? (press Enter to skip): ").strip()  # noqa: ASYNC250
+    if not selection:
+        print("No resume selected. Exiting.")
+        return
+
+    try:
+        idx = int(selection)
+    except ValueError:
+        print("Invalid input; exiting.")
+        return
+
+    if not 0 <= idx < len(sorted_cps):
+        print("Index out of range; exiting.")
+        return
+
+    chosen = sorted_cps[idx]
+
+    new_workflow = create_workflow(checkpoint_storage=storage)
+    # Resume with a fresh workflow instance. The checkpoint carries the
+    # persistent state while this object holds the runtime wiring.
+    result = await run_interactive_session(new_workflow, checkpoint_id=chosen.checkpoint_id)
+    print(f"Workflow completed with: {result}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/03-workflows/checkpoint/checkpoint_with_resume.py

@@ -0,0 +1,157 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Sample: Checkpointing and Resuming a Workflow
+
+Purpose:
+This sample shows how to enable checkpointing for a long-running workflow
+that can be paused and resumed.
+
+What you learn:
+- How to configure checkpointing storage (InMemoryCheckpointStorage for testing)
+- How to resume a workflow from a checkpoint after interruption
+- How to implement executor state management with checkpoint hooks
+- How to handle workflow interruptions and automatic recovery
+
+Pipeline:
+This sample shows a workflow that computes factor pairs for numbers up to a given limit:
+1) A start executor that receives the upper limit and creates the initial task
+2) A worker executor that processes each number to find its factor pairs
+3) The worker uses checkpoint hooks to save/restore its internal state
+
+Prerequisites:
+- Basic understanding of workflow concepts, including executors, edges, events, etc.
+"""
+
+import asyncio
+import sys
+from dataclasses import dataclass
+from random import random
+from typing import Any
+
+from agent_framework import (
+    Executor,
+    InMemoryCheckpointStorage,
+    WorkflowBuilder,
+    WorkflowCheckpoint,
+    WorkflowContext,
+    handler,
+)
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+
+@dataclass
+class ComputeTask:
+    """Task containing the list of numbers remaining to be processed."""
+
+    remaining_numbers: list[int]
+
+
+class StartExecutor(Executor):
+    """Initiates the workflow by providing the upper limit for factor pair computation."""
+
+    @handler
+    async def start(self, upper_limit: int, ctx: WorkflowContext[ComputeTask]) -> None:
+        """Start the workflow with a list of numbers to process."""
+        print(f"StartExecutor: Starting factor pair computation up to {upper_limit}")
+        await ctx.send_message(ComputeTask(remaining_numbers=list(range(1, upper_limit + 1))))
+
+
+class WorkerExecutor(Executor):
+    """Processes numbers to compute their factor pairs and manages executor state for checkpointing."""
+
+    def __init__(self, id: str) -> None:
+        super().__init__(id=id)
+        self._composite_number_pairs: dict[int, list[tuple[int, int]]] = {}
+
+    @handler
+    async def compute(
+        self,
+        task: ComputeTask,
+        ctx: WorkflowContext[ComputeTask, dict[int, list[tuple[int, int]]]],
+    ) -> None:
+        """Process the next number in the task, computing its factor pairs."""
+        next_number = task.remaining_numbers.pop(0)
+
+        print(f"WorkerExecutor: Computing factor pairs for {next_number}")
+        pairs: list[tuple[int, int]] = []
+        for i in range(1, next_number):
+            if next_number % i == 0:
+                pairs.append((i, next_number // i))
+        self._composite_number_pairs[next_number] = pairs
+
+        if not task.remaining_numbers:
+            # All numbers processed - output the results
+            await ctx.yield_output(self._composite_number_pairs)
+        else:
+            # More numbers to process - continue with remaining task
+            await ctx.send_message(task)
+
+    @override
+    async def on_checkpoint_save(self) -> dict[str, Any]:
+        """Save the executor's internal state for checkpointing."""
+        return {"composite_number_pairs": self._composite_number_pairs}
+
+    @override
+    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+        """Restore the executor's internal state from a checkpoint."""
+        self._composite_number_pairs = state.get("composite_number_pairs", {})
+
+
+async def main():
+    # Build workflow with checkpointing enabled
+    checkpoint_storage = InMemoryCheckpointStorage()
+    start = StartExecutor(id="start")
+    worker = WorkerExecutor(id="worker")
+    workflow_builder = (
+        WorkflowBuilder(start_executor=start, checkpoint_storage=checkpoint_storage)
+        .add_edge(start, worker)
+        .add_edge(worker, worker)  # Self-loop for iterative processing
+    )
+
+    # Run workflow with automatic checkpoint recovery
+    latest_checkpoint: WorkflowCheckpoint | None = None
+    while True:
+        workflow = workflow_builder.build()
+
+        # Start from checkpoint or fresh execution
+        print(f"\n** Workflow {workflow.id} started **")
+        event_stream = (
+            workflow.run(message=10, stream=True)
+            if latest_checkpoint is None
+            else workflow.run(checkpoint_id=latest_checkpoint.checkpoint_id, stream=True)
+        )
+
+        output: str | None = None
+        async for event in event_stream:
+            if event.type == "output":
+                output = event.data
+                break
+            if event.type == "superstep_completed" and random() < 0.5:
+                # Randomly simulate system interruptions
+                # The type="superstep_completed" event ensures we only interrupt after
+                # the current super-step is fully complete and checkpointed.
+                # If we interrupt mid-step, the workflow may resume from an earlier point.
+                print("\n** Simulating workflow interruption. Stopping execution. **")
+                break
+
+        # Find the latest checkpoint to resume from
+        latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=workflow.name)
+        if not latest_checkpoint:
+            raise RuntimeError("No checkpoints available to resume from.")
+        print(
+            f"Checkpoint {latest_checkpoint.checkpoint_id}: "
+            f"(iter={latest_checkpoint.iteration_count}, messages={latest_checkpoint.messages})"
+        )
+
+        if output is not None:
+            print(f"\nWorkflow completed successfully with output: {output}")
+            break
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py

@@ -0,0 +1,201 @@
+# Copyright (c) Microsoft. All rights reserved.
+# ruff: noqa: T201
+
+"""Sample: Workflow Checkpointing with Cosmos DB NoSQL.
+
+Purpose:
+This sample shows how to use Azure Cosmos DB NoSQL as a persistent checkpoint
+storage backend for workflows, enabling durable pause-and-resume across
+process restarts.
+
+What you learn:
+- How to configure CosmosCheckpointStorage for workflow checkpointing
+- How to run a workflow that automatically persists checkpoints to Cosmos DB
+- How to resume a workflow from a Cosmos DB checkpoint
+- How to list and inspect available checkpoints
+
+Prerequisites:
+- An Azure Cosmos DB account (or local emulator)
+- Environment variables set (see below)
+
+Environment variables:
+  AZURE_COSMOS_ENDPOINT            - Cosmos DB account endpoint
+  AZURE_COSMOS_DATABASE_NAME       - Database name
+  AZURE_COSMOS_CONTAINER_NAME      - Container name for checkpoints
+Optional:
+  AZURE_COSMOS_KEY                 - Account key (if not using Azure credentials)
+"""
+
+import asyncio
+import os
+import sys
+from dataclasses import dataclass
+from typing import Any
+
+from agent_framework import (
+    Executor,
+    WorkflowBuilder,
+    WorkflowCheckpoint,
+    WorkflowContext,
+    handler,
+)
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+from agent_framework_azure_cosmos import CosmosCheckpointStorage
+
+
+@dataclass
+class ComputeTask:
+    """Task containing the list of numbers remaining to be processed."""
+
+    remaining_numbers: list[int]
+
+
+class StartExecutor(Executor):
+    """Initiates the workflow by providing the upper limit."""
+
+    @handler
+    async def start(self, upper_limit: int, ctx: WorkflowContext[ComputeTask]) -> None:
+        """Start the workflow with numbers up to the given limit."""
+        print(f"StartExecutor: Starting computation up to {upper_limit}")
+        await ctx.send_message(ComputeTask(remaining_numbers=list(range(1, upper_limit + 1))))
+
+
+class WorkerExecutor(Executor):
+    """Processes numbers and manages executor state for checkpointing."""
+
+    def __init__(self, id: str) -> None:
+        """Initialize the worker executor."""
+        super().__init__(id=id)
+        self._results: dict[int, list[tuple[int, int]]] = {}
+
+    @handler
+    async def compute(
+        self,
+        task: ComputeTask,
+        ctx: WorkflowContext[ComputeTask, dict[int, list[tuple[int, int]]]],
+    ) -> None:
+        """Process the next number, computing its factor pairs."""
+        next_number = task.remaining_numbers.pop(0)
+        print(f"WorkerExecutor: Processing {next_number}")
+
+        pairs: list[tuple[int, int]] = []
+        for i in range(1, next_number):
+            if next_number % i == 0:
+                pairs.append((i, next_number // i))
+        self._results[next_number] = pairs
+
+        if not task.remaining_numbers:
+            await ctx.yield_output(self._results)
+        else:
+            await ctx.send_message(task)
+
+    @override
+    async def on_checkpoint_save(self) -> dict[str, Any]:
+        return {"results": self._results}
+
+    @override
+    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+        self._results = state.get("results", {})
+
+
+async def main() -> None:
+    """Run the workflow checkpointing sample with Cosmos DB."""
+    cosmos_endpoint = os.getenv("AZURE_COSMOS_ENDPOINT")
+    cosmos_database_name = os.getenv("AZURE_COSMOS_DATABASE_NAME")
+    cosmos_container_name = os.getenv("AZURE_COSMOS_CONTAINER_NAME")
+    cosmos_key = os.getenv("AZURE_COSMOS_KEY")
+
+    if not cosmos_endpoint or not cosmos_database_name or not cosmos_container_name:
+        print("Please set AZURE_COSMOS_ENDPOINT, AZURE_COSMOS_DATABASE_NAME, and AZURE_COSMOS_CONTAINER_NAME.")
+        return
+
+    # Authentication: supports both managed identity/RBAC and key-based auth.
+    # When AZURE_COSMOS_KEY is set, key-based auth is used.
+    # Otherwise, falls back to DefaultAzureCredential (properly closed via async with).
+    if cosmos_key:
+        async with CosmosCheckpointStorage(
+            endpoint=cosmos_endpoint,
+            credential=cosmos_key,
+            database_name=cosmos_database_name,
+            container_name=cosmos_container_name,
+        ) as checkpoint_storage:
+            await _run_workflow(checkpoint_storage)
+    else:
+        from azure.identity.aio import DefaultAzureCredential
+
+        async with (
+            DefaultAzureCredential() as credential,
+            CosmosCheckpointStorage(
+                endpoint=cosmos_endpoint,
+                credential=credential,
+                database_name=cosmos_database_name,
+                container_name=cosmos_container_name,
+            ) as checkpoint_storage,
+        ):
+            await _run_workflow(checkpoint_storage)
+
+
+async def _run_workflow(checkpoint_storage: CosmosCheckpointStorage) -> None:
+    """Build and run the workflow with Cosmos DB checkpointing."""
+    start = StartExecutor(id="start")
+    worker = WorkerExecutor(id="worker")
+    workflow_builder = (
+        WorkflowBuilder(start_executor=start, checkpoint_storage=checkpoint_storage)
+        .add_edge(start, worker)
+        .add_edge(worker, worker)
+    )
+
+    # --- First run: execute the workflow ---
+    print("\n=== First Run ===\n")
+    workflow = workflow_builder.build()
+
+    output = None
+    async for event in workflow.run(message=8, stream=True):
+        if event.type == "output":
+            output = event.data
+
+    print(f"Factor pairs computed: {output}")
+
+    # List checkpoints saved in Cosmos DB
+    checkpoint_ids = await checkpoint_storage.list_checkpoint_ids(
+        workflow_name=workflow.name,
+    )
+    print(f"\nCheckpoints in Cosmos DB: {len(checkpoint_ids)}")
+    for cid in checkpoint_ids:
+        print(f"  - {cid}")
+
+    # Get the latest checkpoint
+    latest: WorkflowCheckpoint | None = await checkpoint_storage.get_latest(
+        workflow_name=workflow.name,
+    )
+
+    if latest is None:
+        print("No checkpoint found to resume from.")
+        return
+
+    print(f"\nLatest checkpoint: {latest.checkpoint_id}")
+    print(f"  iteration_count: {latest.iteration_count}")
+    print(f"  timestamp: {latest.timestamp}")
+
+    # --- Second run: resume from the latest checkpoint ---
+    print("\n=== Resuming from Checkpoint ===\n")
+    workflow2 = workflow_builder.build()
+
+    output2 = None
+    async for event in workflow2.run(checkpoint_id=latest.checkpoint_id, stream=True):
+        if event.type == "output":
+            output2 = event.data
+
+    if output2:
+        print(f"Resumed workflow produced: {output2}")
+    else:
+        print("Resumed workflow completed (no remaining work — already finished).")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing_foundry.py

@@ -0,0 +1,141 @@
+# Copyright (c) Microsoft. All rights reserved.
+# ruff: noqa: T201
+
+"""Sample: Workflow Checkpointing with Cosmos DB and Azure AI Foundry.
+
+Purpose:
+This sample demonstrates how to use CosmosCheckpointStorage with agents built
+on Azure AI Foundry (via FoundryChatClient). It shows a multi-agent
+workflow where checkpoint state is persisted to Cosmos DB, enabling durable
+pause-and-resume across process restarts.
+
+What you learn:
+- How to wire CosmosCheckpointStorage with FoundryChatClient agents
+- How to combine session history with workflow checkpointing
+- How to resume a workflow-as-agent from a Cosmos DB checkpoint
+
+Key concepts:
+- AgentSession: Maintains conversation history across agent invocations
+- CosmosCheckpointStorage: Persists workflow execution state in Cosmos DB
+- These are complementary: sessions track conversation, checkpoints track workflow state
+
+Environment variables:
+  FOUNDRY_PROJECT_ENDPOINT                - Azure AI Foundry project endpoint
+  FOUNDRY_MODEL                           - Model deployment name
+  AZURE_COSMOS_ENDPOINT                  - Cosmos DB account endpoint
+  AZURE_COSMOS_DATABASE_NAME             - Database name
+  AZURE_COSMOS_CONTAINER_NAME            - Container name for checkpoints
+Optional:
+  AZURE_COSMOS_KEY                       - Account key (if not using Azure credentials)
+"""
+
+import asyncio
+import os
+from typing import Any
+
+from agent_framework import Agent
+from agent_framework.foundry import FoundryChatClient
+from agent_framework.orchestrations import SequentialBuilder
+from agent_framework_azure_cosmos import CosmosCheckpointStorage
+from azure.identity.aio import AzureCliCredential
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+async def main() -> None:
+    """Run the Azure AI Foundry + Cosmos DB checkpointing sample."""
+    project_endpoint = os.getenv("FOUNDRY_PROJECT_ENDPOINT")
+    model = os.getenv("FOUNDRY_MODEL")
+    cosmos_endpoint = os.getenv("AZURE_COSMOS_ENDPOINT")
+    cosmos_database_name = os.getenv("AZURE_COSMOS_DATABASE_NAME")
+    cosmos_container_name = os.getenv("AZURE_COSMOS_CONTAINER_NAME")
+    cosmos_key = os.getenv("AZURE_COSMOS_KEY")
+
+    if not project_endpoint or not model:
+        print("Please set FOUNDRY_PROJECT_ENDPOINT and FOUNDRY_MODEL.")
+        return
+
+    if not cosmos_endpoint or not cosmos_database_name or not cosmos_container_name:
+        print("Please set AZURE_COSMOS_ENDPOINT, AZURE_COSMOS_DATABASE_NAME, and AZURE_COSMOS_CONTAINER_NAME.")
+        return
+
+    # Use a single AzureCliCredential for both Cosmos and Foundry,
+    # properly closed via async context manager.
+    async with AzureCliCredential() as azure_credential:
+        cosmos_credential: Any = cosmos_key if cosmos_key else azure_credential
+
+        async with CosmosCheckpointStorage(
+            endpoint=cosmos_endpoint,
+            credential=cosmos_credential,
+            database_name=cosmos_database_name,
+            container_name=cosmos_container_name,
+        ) as checkpoint_storage:
+            # Create Azure AI Foundry agents
+            client = FoundryChatClient(
+                project_endpoint=project_endpoint,
+                model=model,
+                credential=azure_credential,
+            )
+
+            assistant = Agent(
+                name="assistant",
+                instructions="You are a helpful assistant. Keep responses brief.",
+                client=client,
+            )
+
+            reviewer = Agent(
+                name="reviewer",
+                instructions="You are a reviewer. Provide a one-sentence summary of the assistant's response.",
+                client=client,
+            )
+
+            # Build a sequential workflow and wrap it as an agent
+            workflow = SequentialBuilder(participants=[assistant, reviewer]).build()
+            agent = workflow.as_agent(name="FoundryCheckpointedAgent")
+
+            # --- First run: execute with Cosmos DB checkpointing ---
+            print("=== First Run ===\n")
+
+            session = agent.create_session()
+            query = "What are the benefits of renewable energy?"
+            print(f"User: {query}")
+
+            response = await agent.run(query, session=session, checkpoint_storage=checkpoint_storage)
+
+            for msg in response.messages:
+                speaker = msg.author_name or msg.role
+                print(f"[{speaker}]: {msg.text}")
+
+            # Show checkpoints persisted in Cosmos DB
+            checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow.name)
+            print(f"\nCheckpoints in Cosmos DB: {len(checkpoints)}")
+            for i, cp in enumerate(checkpoints[:5], 1):
+                print(f"  {i}. {cp.checkpoint_id} (iteration={cp.iteration_count})")
+
+            # --- Second run: continue conversation with checkpoint history ---
+            print("\n=== Second Run (continuing conversation) ===\n")
+
+            query2 = "Can you elaborate on the economic benefits?"
+            print(f"User: {query2}")
+
+            response2 = await agent.run(query2, session=session, checkpoint_storage=checkpoint_storage)
+
+            for msg in response2.messages:
+                speaker = msg.author_name or msg.role
+                print(f"[{speaker}]: {msg.text}")
+
+            # Show total checkpoints
+            all_checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow.name)
+            print(f"\nTotal checkpoints after two runs: {len(all_checkpoints)}")
+
+            # Get latest checkpoint
+            latest = await checkpoint_storage.get_latest(workflow_name=workflow.name)
+            if latest:
+                print(f"Latest checkpoint: {latest.checkpoint_id}")
+                print(f"  iteration_count: {latest.iteration_count}")
+                print(f"  timestamp: {latest.timestamp}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/03-workflows/checkpoint/sub_workflow_checkpoint.py

@@ -0,0 +1,415 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import contextlib
+import json
+import sys
+import uuid
+from dataclasses import dataclass, field, replace
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+from agent_framework import (
+    Executor,
+    FileCheckpointStorage,
+    SubWorkflowRequestMessage,
+    SubWorkflowResponseMessage,
+    Workflow,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    WorkflowExecutor,
+    WorkflowRunState,
+    handler,
+    response_handler,
+)
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+CHECKPOINT_DIR = Path(__file__).with_suffix("").parent / "tmp" / "sub_workflow_checkpoints"
+
+"""
+Sample: Checkpointing for workflows that embed sub-workflows.
+
+This sample shows how a parent workflow that wraps a sub-workflow can:
+- run until the sub-workflow emits a human approval request
+- persist a checkpoint that captures the pending request (including complex payloads)
+- resume later, supplying the human decision directly at restore time
+
+It is intentionally similar in spirit to the orchestration checkpoint sample but
+uses ``WorkflowExecutor`` so we exercise the full parent/sub-workflow round-trip.
+"""
+
+
+def _utc_now() -> datetime:
+    return datetime.now()
+
+
+# ---------------------------------------------------------------------------
+# Messages exchanged inside the sub-workflow
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DraftTask:
+    """Task handed from the parent to the sub-workflow writer."""
+
+    topic: str
+    due: datetime
+    iteration: int = 1
+
+
+@dataclass
+class DraftPackage:
+    """Intermediate draft produced by the sub-workflow writer."""
+
+    topic: str
+    content: str
+    iteration: int
+    created_at: datetime = field(default_factory=_utc_now)
+
+
+@dataclass
+class FinalDraft:
+    """Final deliverable returned to the parent workflow."""
+
+    topic: str
+    content: str
+    iterations: int
+    approved_at: datetime
+
+
+@dataclass
+class ReviewRequest:
+    """Human approval request surfaced via `request_info`."""
+
+    id: str = str(uuid.uuid4())
+    topic: str = ""
+    iteration: int = 1
+    draft_excerpt: str = ""
+    due_iso: str = ""
+    reviewer_guidance: list[str] = field(default_factory=list)  # type: ignore
+
+
+@dataclass
+class ReviewDecision:
+    """The review decision to be sent to downstream executors along with the original request."""
+
+    decision: str
+    original_request: ReviewRequest
+
+
+# ---------------------------------------------------------------------------
+# Sub-workflow executors
+# ---------------------------------------------------------------------------
+
+
+class DraftWriter(Executor):
+    """Produces an initial draft for the supplied topic."""
+
+    def __init__(self) -> None:
+        super().__init__(id="draft_writer")
+
+    @handler
+    async def create_draft(self, task: DraftTask, ctx: WorkflowContext[DraftPackage]) -> None:
+        draft = DraftPackage(
+            topic=task.topic,
+            content=(
+                f"Launch plan for {task.topic}.\n\n"
+                "- Outline the customer message.\n"
+                "- Highlight three differentiators.\n"
+                "- Close with a next-step CTA.\n"
+                f"(iteration {task.iteration})"
+            ),
+            iteration=task.iteration,
+        )
+        await ctx.send_message(draft, target_id="draft_review")
+
+
+class DraftReviewRouter(Executor):
+    """Turns draft packages into human approval requests."""
+
+    def __init__(self) -> None:
+        super().__init__(id="draft_review")
+
+    @handler
+    async def request_review(self, draft: DraftPackage, ctx: WorkflowContext) -> None:
+        """Request a review upon receiving a draft."""
+        excerpt = draft.content.splitlines()[0]
+        request = ReviewRequest(
+            topic=draft.topic,
+            iteration=draft.iteration,
+            draft_excerpt=excerpt,
+            due_iso=draft.created_at.isoformat(),
+            reviewer_guidance=[
+                "Ensure tone matches launch messaging",
+                "Confirm CTA is action-oriented",
+            ],
+        )
+        await ctx.request_info(request_data=request, response_type=str)
+
+    @response_handler
+    async def forward_decision(
+        self,
+        original_request: ReviewRequest,
+        decision: str,
+        ctx: WorkflowContext[ReviewDecision],
+    ) -> None:
+        """Route the decision to the next executor."""
+        await ctx.send_message(ReviewDecision(decision=decision, original_request=original_request))
+
+
+class DraftFinaliser(Executor):
+    """Applies the human decision and emits the final draft."""
+
+    def __init__(self) -> None:
+        super().__init__(id="draft_finaliser")
+
+    @handler
+    async def on_review_decision(
+        self,
+        review_decision: ReviewDecision,
+        ctx: WorkflowContext[DraftTask, FinalDraft],
+    ) -> None:
+        reply = review_decision.decision.strip().lower()
+        original = review_decision.original_request
+        topic = original.topic if original else "unknown topic"
+        iteration = original.iteration if original else 1
+
+        if reply != "approve":
+            # Loop back with a follow-up task. In a real workflow you would
+            # incorporate the human guidance; here we just increment the counter.
+            next_task = DraftTask(
+                topic=topic,
+                due=_utc_now() + timedelta(hours=1),
+                iteration=iteration + 1,
+            )
+            await ctx.send_message(next_task, target_id="draft_writer")
+            return
+
+        final = FinalDraft(
+            topic=topic,
+            content=f"Approved launch narrative for {topic} (iteration {iteration}).",
+            iterations=iteration,
+            approved_at=_utc_now(),
+        )
+        await ctx.yield_output(final)
+
+
+# ---------------------------------------------------------------------------
+# Parent workflow executors
+# ---------------------------------------------------------------------------
+
+
+class LaunchCoordinator(Executor):
+    """Owns the top-level workflow and collects the final draft."""
+
+    def __init__(self) -> None:
+        super().__init__(id="launch_coordinator")
+        # Track pending requests to match responses
+        self._pending_requests: dict[str, SubWorkflowRequestMessage] = {}
+
+    @handler
+    async def kick_off(self, topic: str, ctx: WorkflowContext[DraftTask]) -> None:
+        task = DraftTask(topic=topic, due=_utc_now() + timedelta(hours=2))
+        await ctx.send_message(task)
+
+    @handler
+    async def collect_final(self, draft: FinalDraft, ctx: WorkflowContext[None, FinalDraft]) -> None:
+        approved_at = draft.approved_at
+        normalised = draft
+        if isinstance(approved_at, str):
+            with contextlib.suppress(ValueError):
+                parsed = datetime.fromisoformat(approved_at)
+                normalised = replace(draft, approved_at=parsed)
+                approved_at = parsed
+
+        approved_display = approved_at.isoformat() if hasattr(approved_at, "isoformat") else str(approved_at)
+
+        print("\n>>> Parent workflow received approved draft:")
+        print(f"- Topic: {normalised.topic}")
+        print(f"- Iterations: {normalised.iterations}")
+        print(f"- Approved at: {approved_display}")
+        print(f"- Content: {normalised.content}\n")
+
+        await ctx.yield_output(normalised)
+
+    @handler
+    async def handler_sub_workflow_request(
+        self,
+        request: SubWorkflowRequestMessage,
+        ctx: WorkflowContext,
+    ) -> None:
+        """Handle requests from the sub-workflow.
+
+        Note that the message type must be SubWorkflowRequestMessage to intercept the request.
+        """
+        if not isinstance(request.source_event.data, ReviewRequest):
+            raise TypeError(f"Expected 'ReviewRequest', got {type(request.source_event.data)}")
+
+        # Record the request for response matching
+        review_request = request.source_event.data
+        self._pending_requests[review_request.id] = request
+
+        # Send the request without modification
+        await ctx.request_info(request_data=review_request, response_type=str)
+
+    @response_handler
+    async def handle_request_response(
+        self,
+        original_request: ReviewRequest,
+        response: str,
+        ctx: WorkflowContext[SubWorkflowResponseMessage],
+    ) -> None:
+        """Process the response and send it back to the sub-workflow.
+
+        Note that the response must be sent back using SubWorkflowResponseMessage to route
+        the response back to the sub-workflow.
+        """
+        request_message = self._pending_requests.pop(original_request.id, None)
+
+        if request_message is None:
+            raise ValueError("No matching pending request found for the resource response")
+
+        await ctx.send_message(request_message.create_response(response))
+
+    @override
+    async def on_checkpoint_save(self) -> dict[str, Any]:
+        """Capture any additional state needed for checkpointing."""
+        return {
+            "pending_requests": self._pending_requests,
+        }
+
+    @override
+    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+        """Restore any additional state needed from checkpointing."""
+        self._pending_requests = state.get("pending_requests", {})
+
+
+# ---------------------------------------------------------------------------
+# Workflow construction helpers
+# ---------------------------------------------------------------------------
+
+
+def build_sub_workflow() -> WorkflowExecutor:
+    """Assemble the sub-workflow used by the parent workflow executor."""
+    writer = DraftWriter()
+    router = DraftReviewRouter()
+    finaliser = DraftFinaliser()
+    sub_workflow = (
+        WorkflowBuilder(start_executor=writer)
+        .add_edge(writer, router)
+        .add_edge(router, finaliser)
+        .add_edge(finaliser, writer)  # permits revision loops
+        .build()
+    )
+
+    return WorkflowExecutor(sub_workflow, id="launch_subworkflow")
+
+
+def build_parent_workflow(storage: FileCheckpointStorage) -> Workflow:
+    """Assemble the parent workflow that embeds the sub-workflow."""
+    coordinator = LaunchCoordinator()
+    sub_executor = build_sub_workflow()
+    return (
+        WorkflowBuilder(start_executor=coordinator, checkpoint_storage=storage)
+        .add_edge(coordinator, sub_executor)
+        .add_edge(sub_executor, coordinator)
+        .build()
+    )
+
+
+async def main() -> None:
+    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
+    for file in CHECKPOINT_DIR.glob("*.json"):
+        file.unlink()
+
+    storage = FileCheckpointStorage(CHECKPOINT_DIR)
+
+    workflow = build_parent_workflow(storage)
+
+    print("\n=== Stage 1: run until sub-workflow requests human review ===")
+
+    request_id: str | None = None
+    async for event in workflow.run("Contoso Gadget Launch", stream=True):
+        if event.type == "request_info" and request_id is None:
+            request_id = event.request_id
+            print(f"Captured review request id: {request_id}")
+        if event.type == "status" and event.state is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
+            break
+
+    if request_id is None:
+        raise RuntimeError("Sub-workflow completed without requesting review.")
+
+    resume_checkpoint = await storage.get_latest(workflow_name=workflow.name)
+    if not resume_checkpoint:
+        raise RuntimeError("No checkpoints found.")
+
+    # Print the checkpoint to show pending requests
+    # We didn't handle the request above so the request is still pending the last checkpoint
+    print(f"Using checkpoint {resume_checkpoint.checkpoint_id} at iteration {resume_checkpoint.iteration_count}")
+
+    checkpoint_path = storage.storage_path / f"{resume_checkpoint.checkpoint_id}.json"
+    if checkpoint_path.exists():
+        checkpoint_content_dict = json.loads(checkpoint_path.read_text())
+        print(f"Pending review requests: {checkpoint_content_dict.get('pending_request_info_events', {})}")
+
+    print("\n=== Stage 2: resume from checkpoint ===")
+
+    # Rebuild fresh instances to mimic a separate process resuming
+    workflow2 = build_parent_workflow(storage)
+
+    request_info_event: WorkflowEvent | None = None
+    async for event in workflow2.run(checkpoint_id=resume_checkpoint.checkpoint_id, stream=True):
+        if event.type == "request_info":
+            request_info_event = event
+
+    if request_info_event is None:
+        raise RuntimeError("No request_info_event captured.")
+
+    print("\n=== Stage 3: approve draft ==")
+
+    approval_response = "approve"
+    output_event: WorkflowEvent | None = None
+    async for event in workflow2.run(stream=True, responses={request_info_event.request_id: approval_response}):
+        if event.type == "output":
+            output_event = event
+
+    if output_event is None:
+        raise RuntimeError("Workflow did not complete after resume.")
+
+    output = output_event.data
+    print("\n=== Final Draft (from resumed run) ===")
+    print(output)
+
+    """"
+    Sample Output:
+
+    === Stage 1: run until sub-workflow requests human review ===
+    Captured review request id: 032c9f3a-ad1b-4a52-89be-a168d6663011
+    Using checkpoint 54f376c2-f849-44e4-9d8d-e627fd27ab96 at iteration 2
+    Pending review requests (sub executor snapshot): []
+    Pending review requests (parent executor snapshot): ['032c9f3a-ad1b-4a52-89be-a168d6663011']
+
+    === Stage 2: resume from checkpoint and approve draft ===
+
+    >>> Parent workflow received approved draft:
+    - Topic: Contoso Gadget Launch
+    - Iterations: 1
+    - Approved at: 2025-09-25T14:29:34.479164
+    - Content: Approved launch narrative for Contoso Gadget Launch (iteration 1).
+
+
+    === Final Draft (from resumed run) ===
+    FinalDraft(topic='Contoso Gadget Launch', content='Approved launch narrative for Contoso
+    Gadget Launch (iteration 1).', iterations=1, approved_at=datetime.datetime(2025, 9, 25, 14, 29, 34, 479164))
+    Coordinator stored final draft successfully.
+    """
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/03-workflows/checkpoint/workflow_as_agent_checkpoint.py

@@ -0,0 +1,186 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Sample: Workflow as Agent with Checkpointing
+
+Purpose:
+This sample demonstrates how to use checkpointing with a workflow wrapped as an agent.
+It shows how to enable checkpoint storage when calling agent.run(),
+allowing workflow execution state to be persisted and potentially resumed.
+
+What you learn:
+- How to pass checkpoint_storage to WorkflowAgent.run()
+- How checkpoints are created during workflow-as-agent execution
+- How to combine thread conversation history with workflow checkpointing
+- How to resume a workflow-as-agent from a checkpoint
+
+Key concepts:
+- Thread (AgentSession): Maintains conversation history across agent invocations
+- Checkpoint: Persists workflow execution state for pause/resume capability
+- These are complementary: sessions track conversation, checkpoints track workflow state
+
+Prerequisites:
+- FOUNDRY_PROJECT_ENDPOINT must be your Azure AI Foundry Agent Service (V2) project endpoint.
+- FOUNDRY_MODEL must be set to your Azure OpenAI model deployment name.
+"""
+
+import asyncio
+import os
+
+from agent_framework import (
+    Agent,
+    InMemoryCheckpointStorage,
+    InMemoryHistoryProvider,
+)
+from agent_framework.foundry import FoundryChatClient
+from agent_framework.orchestrations import SequentialBuilder
+from azure.identity import AzureCliCredential
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+
+async def basic_checkpointing() -> None:
+    """Demonstrate basic checkpoint storage with workflow-as-agent."""
+
+    print("=" * 60)
+    print("Basic Checkpointing with Workflow as Agent")
+    print("=" * 60)
+
+    client = FoundryChatClient(
+        project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+        model=os.environ["FOUNDRY_MODEL"],
+        credential=AzureCliCredential(),
+    )
+
+    assistant = Agent(
+        client=client,
+        name="assistant",
+        instructions="You are a helpful assistant. Keep responses brief.",
+    )
+
+    reviewer = Agent(
+        client=client,
+        name="reviewer",
+        instructions="You are a reviewer. Provide a one-sentence summary of the assistant's response.",
+    )
+
+    workflow = SequentialBuilder(participants=[assistant, reviewer]).build()
+    agent = workflow.as_agent()
+
+    # Create checkpoint storage
+    checkpoint_storage = InMemoryCheckpointStorage()
+
+    # Run with checkpointing enabled
+    query = "What are the benefits of renewable energy?"
+    print(f"\nUser: {query}")
+
+    response = await agent.run(query, checkpoint_storage=checkpoint_storage)
+
+    for msg in response.messages:
+        speaker = msg.author_name or msg.role
+        print(f"[{speaker}]: {msg.text}")
+
+    # Show checkpoints that were created
+    checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow.name)
+    print(f"\nCheckpoints created: {len(checkpoints)}")
+    for i, cp in enumerate(checkpoints[:5], 1):
+        print(f"  {i}. {cp.checkpoint_id}")
+
+
+async def checkpointing_with_thread() -> None:
+    """Demonstrate combining thread history with checkpointing."""
+    print("\n" + "=" * 60)
+    print("Checkpointing with Thread Conversation History")
+    print("=" * 60)
+
+    client = FoundryChatClient(
+        project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+        model=os.environ["FOUNDRY_MODEL"],
+        credential=AzureCliCredential(),
+    )
+
+    assistant = Agent(
+        client=client,
+        name="memory_assistant",
+        instructions="You are a helpful assistant with good memory. Reference previous conversation when relevant.",
+    )
+
+    workflow = SequentialBuilder(participants=[assistant]).build()
+    agent = workflow.as_agent()
+
+    # Create both session (for conversation) and checkpoint storage (for workflow state)
+    session = agent.create_session()
+    checkpoint_storage = InMemoryCheckpointStorage()
+
+    # First turn
+    query1 = "My favorite color is blue. Remember that."
+    print(f"\n[Turn 1] User: {query1}")
+    response1 = await agent.run(query1, session=session, checkpoint_storage=checkpoint_storage)
+    if response1.messages:
+        print(f"[assistant]: {response1.messages[0].text}")
+
+    # Second turn - agent should remember from session history
+    query2 = "What's my favorite color?"
+    print(f"\n[Turn 2] User: {query2}")
+    response2 = await agent.run(query2, session=session, checkpoint_storage=checkpoint_storage)
+    if response2.messages:
+        print(f"[assistant]: {response2.messages[0].text}")
+
+    # Show accumulated state
+    checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow.name)
+    print(f"\nTotal checkpoints across both turns: {len(checkpoints)}")
+
+    memory_state = session.state.get(InMemoryHistoryProvider.DEFAULT_SOURCE_ID, {})
+    history = memory_state.get("messages", [])
+    print(f"Messages in session history: {len(history)}")
+
+
+async def streaming_with_checkpoints() -> None:
+    """Demonstrate streaming with checkpoint storage."""
+    print("\n" + "=" * 60)
+    print("Streaming with Checkpointing")
+    print("=" * 60)
+
+    client = FoundryChatClient(
+        project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+        model=os.environ["FOUNDRY_MODEL"],
+        credential=AzureCliCredential(),
+    )
+
+    assistant = Agent(
+        client=client,
+        name="streaming_assistant",
+        instructions="You are a helpful assistant.",
+    )
+
+    workflow = SequentialBuilder(participants=[assistant]).build()
+    agent = workflow.as_agent()
+
+    checkpoint_storage = InMemoryCheckpointStorage()
+
+    query = "List three interesting facts about the ocean."
+    print(f"\nUser: {query}")
+    print("[assistant]: ", end="", flush=True)
+
+    # Stream with checkpointing
+    async for update in agent.run(query, checkpoint_storage=checkpoint_storage, stream=True):
+        if update.text:
+            print(update.text, end="", flush=True)
+
+    print()  # Newline after streaming
+
+    checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow.name)
+    print(f"\nCheckpoints created during stream: {len(checkpoints)}")
+
+
+async def main() -> None:
+    """Run all checkpoint examples."""
+    await basic_checkpointing()
+    await checkpointing_with_thread()
+    await streaming_with_checkpoints()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())