chore: import upstream snapshot with attribution

python/packages/durabletask/tests/integration_tests/.env.example

@@ -0,0 +1,13 @@
+# Azure OpenAI Configuration
+AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
+AZURE_OPENAI_MODEL=your-deployment-name
+# Optional: Use Azure CLI authentication if not provided
+# AZURE_OPENAI_API_KEY=your-api-key
+
+# Durable Task Scheduler Configuration
+ENDPOINT=http://localhost:8080
+TASKHUB=default
+
+# Redis Configuration (for streaming tests)
+REDIS_CONNECTION_STRING=redis://localhost:6379
+REDIS_STREAM_TTL_MINUTES=10

python/packages/durabletask/tests/integration_tests/README.md

@@ -0,0 +1,110 @@
+# Sample Integration Tests
+
+Integration tests that validate the Durable Agent Framework samples by running them against a Durable Task Scheduler (DTS) instance.
+
+## Setup
+
+### 1. Create `.env` file
+
+Copy `.env.example` to `.env` and fill in your Azure credentials:
+
+```bash
+cp .env.example .env
+```
+
+Required variables:
+- `AZURE_OPENAI_ENDPOINT`
+- `AZURE_OPENAI_MODEL`
+- `AZURE_OPENAI_API_KEY` (optional if using Azure CLI authentication)
+- `ENDPOINT` (default: http://localhost:8080)
+- `TASKHUB` (default: default)
+
+Optional variables (for streaming tests):
+- `REDIS_CONNECTION_STRING` (default: redis://localhost:6379)
+- `REDIS_STREAM_TTL_MINUTES` (default: 10)
+
+### 2. Start required services
+
+**Durable Task Scheduler:**
+```bash
+docker run -d --name dts-emulator -p 8080:8080 -p 8082:8082 mcr.microsoft.com/dts/dts-emulator:latest
+```
+- Port 8080: gRPC endpoint (used by tests)
+- Port 8082: Web dashboard (optional, for monitoring)
+
+**Redis (for streaming tests):**
+```bash
+docker run -d --name redis -p 6379:6379 redis:latest
+```
+- Port 6379: Redis server endpoint
+
+## Running Tests
+
+The tests automatically start and stop worker processes for each sample.
+
+### Run all sample tests
+```bash
+uv run pytest packages/durabletask/tests/integration_tests -v
+```
+
+### Run specific sample
+```bash
+uv run pytest packages/durabletask/tests/integration_tests/test_01_single_agent.py -v
+```
+
+### Run with verbose output
+```bash
+uv run pytest packages/durabletask/tests/integration_tests -sv
+```
+
+## How It Works
+
+Each test file uses pytest markers to automatically configure and start the worker process:
+
+```python
+pytestmark = [
+    pytest.mark.sample("03_single_agent_streaming"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_azure_openai,
+    pytest.mark.requires_dts,
+    pytest.mark.requires_redis,
+]
+```
+
+## Troubleshooting
+
+**Tests are skipped:**
+Ensure the required environment variables (e.g., `AZURE_OPENAI_ENDPOINT`) are set in your `.env` file.
+
+**DTS connection failed:**
+Check that the DTS emulator container is running: `docker ps | grep dts-emulator`
+
+**Redis connection failed:**
+Check that Redis is running: `docker ps | grep redis`
+
+**Missing environment variables:**
+Ensure your `.env` file contains all required variables from `.env.example`.
+
+**Tests timeout:**
+Check that Azure OpenAI credentials are valid and the service is accessible.
+
+If you see "DTS emulator is not available":
+- Ensure Docker container is running: `docker ps | grep dts-emulator`
+- Check port 8080 is not in use by another process
+- Restart the container if needed
+
+### Azure OpenAI Errors
+
+If you see authentication or deployment errors:
+- Verify your `AZURE_OPENAI_ENDPOINT` is correct
+- Confirm `AZURE_OPENAI_MODEL` matches your deployment
+- If using API key, check `AZURE_OPENAI_API_KEY` is valid
+- If using Azure CLI, ensure you're logged in: `az login`
+
+## CI/CD
+
+For automated testing in CI/CD pipelines:
+
+1. Use Docker Compose to start DTS emulator
+2. Set environment variables via CI/CD secrets
+3. Run tests with appropriate markers: `pytest -m integration_test`

python/packages/durabletask/tests/integration_tests/conftest.py

@@ -0,0 +1,512 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""Pytest configuration and fixtures for durabletask integration tests."""
+
+import asyncio
+import json
+import logging
+import os
+import socket
+import subprocess
+import sys
+import time
+import uuid
+from collections.abc import Generator
+from pathlib import Path
+from typing import Any, Protocol, cast
+from urllib.parse import urlparse
+
+import pytest
+import redis.asyncio as aioredis
+from dotenv import load_dotenv
+from durabletask.azuremanaged.client import DurableTaskSchedulerClient
+from durabletask.client import OrchestrationStatus
+
+from agent_framework_durabletask import DurableAIAgentClient, DurableWorkflowClient
+
+# Load environment variables from .env file
+load_dotenv(Path(__file__).parent / ".env")
+
+# Configure logging to reduce noise during tests
+logging.basicConfig(level=logging.WARNING)
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[DurableTaskSchedulerClient, DurableAIAgentClient]: ...
+
+
+# =============================================================================
+# Environment and Service Checks
+# =============================================================================
+
+
+def _get_dts_endpoint() -> str:
+    """Get the DTS endpoint from environment or use default."""
+    return os.getenv("ENDPOINT", "http://localhost:8080")
+
+
+def _check_dts_available(endpoint: str | None = None) -> bool:
+    """Check if DTS emulator is available at the given endpoint."""
+    try:
+        resolved_endpoint: str = _get_dts_endpoint() if endpoint is None else endpoint
+        parsed = urlparse(resolved_endpoint)
+        host = parsed.hostname or "localhost"
+        port = parsed.port or 8080
+
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(2)
+            return sock.connect_ex((host, port)) == 0
+    except Exception:
+        return False
+
+
+def _check_redis_available() -> bool:
+    """Check if Redis is available at the default connection string."""
+    try:
+
+        async def test_connection() -> bool:
+            redis_url = os.getenv("REDIS_CONNECTION_STRING", "redis://localhost:6379")
+            try:
+                client = aioredis.from_url(redis_url, socket_timeout=2)  # type: ignore[reportUnknownMemberType]
+                await client.ping()  # type: ignore[reportUnknownMemberType]
+                await client.aclose()  # type: ignore[reportUnknownMemberType]
+                return True
+            except Exception:
+                return False
+
+        return asyncio.run(test_connection())
+    except Exception:
+        return False
+
+
+# =============================================================================
+# Client Factory Functions
+# =============================================================================
+
+
+def create_dts_client(endpoint: str, taskhub: str) -> DurableTaskSchedulerClient:
+    """Create a DurableTaskSchedulerClient with common configuration.
+
+    Args:
+        endpoint: The DTS endpoint address
+        taskhub: The task hub name
+
+    Returns:
+        A configured DurableTaskSchedulerClient instance
+    """
+    return DurableTaskSchedulerClient(
+        host_address=endpoint,
+        secure_channel=False,
+        taskhub=taskhub,
+        token_credential=None,
+    )
+
+
+def create_agent_client(
+    endpoint: str,
+    taskhub: str,
+    max_poll_retries: int = 90,
+) -> tuple[DurableTaskSchedulerClient, DurableAIAgentClient]:
+    """Create a DurableAIAgentClient with the underlying DTS client.
+
+    Args:
+        endpoint: The DTS endpoint address
+        taskhub: The task hub name
+        max_poll_retries: Max poll retries for the agent client
+
+    Returns:
+        A tuple of (DurableTaskSchedulerClient, DurableAIAgentClient)
+    """
+    dts_client = create_dts_client(endpoint, taskhub)
+    agent_client = DurableAIAgentClient(dts_client, max_poll_retries=max_poll_retries)
+    return dts_client, agent_client
+
+
+# =============================================================================
+# Orchestration Helper Class
+# =============================================================================
+
+
+class OrchestrationHelper:
+    """Helper class for orchestration-related test operations."""
+
+    def __init__(self, dts_client: DurableTaskSchedulerClient):
+        """Initialize the orchestration helper.
+
+        Args:
+            dts_client: The DurableTaskSchedulerClient instance to use
+        """
+        self.client = dts_client
+
+    def wait_for_orchestration(
+        self,
+        instance_id: str,
+        timeout: float = 60.0,
+    ) -> Any:
+        """Wait for an orchestration to complete.
+
+        Args:
+            instance_id: The orchestration instance ID
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            The final OrchestrationMetadata
+
+        Raises:
+            TimeoutError: If the orchestration doesn't complete within timeout
+            RuntimeError: If the orchestration fails
+        """
+        # Use the built-in wait_for_orchestration_completion method
+        metadata = self.client.wait_for_orchestration_completion(
+            instance_id=instance_id,
+            timeout=int(timeout),
+        )
+
+        if metadata is None:
+            raise TimeoutError(f"Orchestration {instance_id} did not complete within {timeout} seconds")
+
+        # Check if failed or terminated
+        if metadata.runtime_status == OrchestrationStatus.FAILED:
+            raise RuntimeError(f"Orchestration {instance_id} failed: {metadata.serialized_custom_status}")
+        if metadata.runtime_status == OrchestrationStatus.TERMINATED:
+            raise RuntimeError(f"Orchestration {instance_id} was terminated")
+
+        return metadata
+
+    def wait_for_orchestration_with_output(
+        self,
+        instance_id: str,
+        timeout: float = 60.0,
+    ) -> tuple[Any, Any]:
+        """Wait for an orchestration to complete and return its output.
+
+        Args:
+            instance_id: The orchestration instance ID
+            timeout: Maximum time to wait in seconds
+
+        Returns:
+            A tuple of (OrchestrationMetadata, output)
+
+        Raises:
+            TimeoutError: If the orchestration doesn't complete within timeout
+            RuntimeError: If the orchestration fails
+        """
+        metadata = self.wait_for_orchestration(instance_id, timeout)
+
+        # The output should be available in the metadata
+        return metadata, metadata.serialized_output
+
+    def get_orchestration_status(self, instance_id: str) -> Any | None:
+        """Get the current status of an orchestration.
+
+        Args:
+            instance_id: The orchestration instance ID
+
+        Returns:
+            The OrchestrationMetadata or None if not found
+        """
+        try:
+            # Try to wait with a short timeout to get current status
+            return self.client.wait_for_orchestration_completion(
+                instance_id=instance_id,
+                timeout=1,  # Very short timeout, just checking status
+            )
+        except Exception:
+            return None
+
+    def raise_event(
+        self,
+        instance_id: str,
+        event_name: str,
+        event_data: Any = None,
+    ) -> None:
+        """Raise an external event to an orchestration.
+
+        Args:
+            instance_id: The orchestration instance ID
+            event_name: The name of the event
+            event_data: The event data payload
+        """
+        self.client.raise_orchestration_event(instance_id, event_name, data=event_data)
+
+    def wait_for_notification(self, instance_id: str, timeout_seconds: int = 30) -> bool:
+        """Wait for the orchestration to reach a notification point.
+
+        Polls the orchestration status until it appears to be waiting for approval.
+
+        Args:
+            instance_id: The orchestration instance ID
+            timeout_seconds: Maximum time to wait
+
+        Returns:
+            True if notification detected, False if timeout
+        """
+        start_time = time.time()
+        while time.time() - start_time < timeout_seconds:
+            try:
+                metadata = self.client.get_orchestration_state(
+                    instance_id=instance_id,
+                )
+
+                if metadata:
+                    # Check if we're waiting for approval by examining custom status
+                    if metadata.serialized_custom_status:
+                        try:
+                            custom_status = json.loads(metadata.serialized_custom_status)
+                            # Handle both string and dict custom status
+                            status_str = custom_status if isinstance(custom_status, str) else str(custom_status)
+                            if status_str.lower().startswith("requesting human feedback"):
+                                return True
+                        except (json.JSONDecodeError, AttributeError):
+                            # If it's not JSON, treat as plain string
+                            if metadata.serialized_custom_status.lower().startswith("requesting human feedback"):
+                                return True
+
+                    # Check for terminal states
+                    if metadata.runtime_status.name == "COMPLETED" or metadata.runtime_status.name == "FAILED":
+                        return False
+            except Exception:
+                # Silently ignore transient errors during polling (e.g., network issues, service unavailable).
+                # The loop will retry until timeout, allowing the service to recover.
+                pass
+
+            time.sleep(1)
+
+        return False
+
+
+# =============================================================================
+# Pytest Configuration
+# =============================================================================
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    """Register custom markers."""
+    config.addinivalue_line("markers", "integration_test: mark test as integration test")
+    config.addinivalue_line("markers", "requires_dts: mark test as requiring DTS emulator")
+    config.addinivalue_line("markers", "requires_azure_openai: mark test as requiring Azure OpenAI")
+    config.addinivalue_line("markers", "requires_redis: mark test as requiring Redis")
+    config.addinivalue_line(
+        "markers",
+        "sample(path): specify the sample directory name for the test (e.g., @pytest.mark.sample('01_single_agent'))",
+    )
+
+
+def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
+    """Skip tests based on markers and environment availability."""
+    foundry_vars = ["FOUNDRY_PROJECT_ENDPOINT", "FOUNDRY_MODEL"]
+    foundry_available = all(os.getenv(var) for var in foundry_vars)
+    azure_openai_vars = ["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_MODEL"]
+    azure_openai_available = all(os.getenv(var) for var in azure_openai_vars)
+    skip_foundry = pytest.mark.skip(reason=f"Missing required environment variables: {', '.join(foundry_vars)}")
+    skip_azure_openai = pytest.mark.skip(
+        reason=f"Missing required environment variables: {', '.join(azure_openai_vars)}"
+    )
+
+    # Check DTS availability
+    dts_available = _check_dts_available()
+    skip_dts = pytest.mark.skip(reason=f"DTS emulator is not available at {_get_dts_endpoint()}")
+
+    # Check Redis availability
+    redis_available = _check_redis_available()
+    skip_redis = pytest.mark.skip(reason="Redis is not available at redis://localhost:6379")
+
+    for item in items:
+        if "requires_azure_openai" in item.keywords and not foundry_available:
+            item.add_marker(skip_foundry)
+        sample_marker = item.get_closest_marker("sample")
+        sample_name = sample_marker.args[0] if sample_marker and sample_marker.args else None
+        if sample_name == "06_multi_agent_orchestration_conditionals" and not azure_openai_available:
+            item.add_marker(skip_azure_openai)
+        if "requires_dts" in item.keywords and not dts_available:
+            item.add_marker(skip_dts)
+        if "requires_redis" in item.keywords and not redis_available:
+            item.add_marker(skip_redis)
+
+
+# =============================================================================
+# Pytest Fixtures
+# =============================================================================
+
+
+@pytest.fixture(scope="session")
+def dts_endpoint() -> str:
+    """Get the DTS endpoint from environment or use default."""
+    return _get_dts_endpoint()
+
+
+@pytest.fixture(scope="session")
+def dts_available(dts_endpoint: str) -> bool:
+    """Check if DTS emulator is available and responding."""
+    if _check_dts_available(dts_endpoint):
+        return True
+    pytest.skip(f"DTS emulator is not available at {dts_endpoint}")
+    return False
+
+
+@pytest.fixture(scope="module")
+def check_sample_env(request: pytest.FixtureRequest) -> None:
+    """Verify the environment variables required by the current sample are set."""
+    sample_marker = request.node.get_closest_marker("sample")  # type: ignore[union-attr]
+    if not sample_marker:
+        pytest.fail("Test class must have @pytest.mark.sample() marker")
+
+    sample_name = cast(str, sample_marker.args[0])  # type: ignore[union-attr]
+    # Samples that host no AI agents need no model credentials (only the DTS emulator).
+    no_llm_samples = {"12_subworkflow_hitl"}
+    if sample_name in no_llm_samples:
+        return
+    if sample_name == "06_multi_agent_orchestration_conditionals":
+        required_vars = ["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_MODEL"]
+    else:
+        required_vars = ["FOUNDRY_PROJECT_ENDPOINT", "FOUNDRY_MODEL"]
+    missing = [var for var in required_vars if not os.getenv(var)]
+
+    if missing:
+        pytest.skip(f"Missing required environment variables: {', '.join(missing)}")
+
+
+@pytest.fixture(scope="module")
+def unique_taskhub() -> str:
+    """Generate a unique task hub name for test isolation."""
+    # Use a shorter UUID to avoid naming issues
+    return f"test-{uuid.uuid4().hex[:8]}"
+
+
+@pytest.fixture(scope="module")
+def worker_process(
+    dts_available: bool,
+    check_sample_env: None,
+    dts_endpoint: str,
+    unique_taskhub: str,
+    request: pytest.FixtureRequest,
+) -> Generator[dict[str, Any], None, None]:
+    """Start a worker process for the current test module by running the sample worker.py.
+
+    This fixture:
+    1. Determines which sample to run from @pytest.mark.sample()
+    2. Starts the sample's worker.py as a subprocess
+    3. Waits for the worker to be ready
+    4. Tears down the worker after tests complete
+
+    Usage:
+    @pytest.mark.sample("01_single_agent")
+    class TestSingleAgent:
+        ...
+    """
+    # Get sample path from marker
+    sample_marker = request.node.get_closest_marker("sample")  # type: ignore[union-attr]
+    if not sample_marker:
+        pytest.fail("Test class must have @pytest.mark.sample() marker")
+
+    sample_name: str = cast(str, sample_marker.args[0])  # type: ignore[union-attr]
+    sample_path: Path = Path(__file__).parents[4] / "samples" / "04-hosting" / "durabletask" / sample_name
+    worker_file: Path = sample_path / "worker.py"
+
+    if not worker_file.exists():
+        pytest.fail(f"Sample worker not found: {worker_file}")
+
+    # Set up environment for worker subprocess
+    env = os.environ.copy()
+    env["ENDPOINT"] = dts_endpoint
+    env["TASKHUB"] = unique_taskhub
+
+    # Start worker subprocess
+    try:
+        # On Windows, use CREATE_NEW_PROCESS_GROUP to allow proper termination
+        # shell=True only on Windows to handle PATH resolution
+        if sys.platform == "win32":
+            process = subprocess.Popen(
+                [sys.executable, str(worker_file)],
+                cwd=str(sample_path),
+                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
+                shell=True,
+                env=env,
+                text=True,
+            )
+        # On Unix, don't use shell=True to avoid shell wrapper issues
+        else:
+            process = subprocess.Popen(
+                [sys.executable, str(worker_file)],
+                cwd=str(sample_path),
+                env=env,
+                text=True,
+            )
+    except Exception as e:
+        pytest.fail(f"Failed to start worker subprocess: {e}")
+
+    # Wait for worker to initialize
+    # The worker needs time to:
+    # 1. Start Python and import modules
+    # 2. Create Azure OpenAI clients
+    # 3. Register agents with the DTS worker
+    # 4. Connect to DTS and be ready to receive signals
+    #
+    # We use a generous wait time because CI environments can be slow,
+    # and the first test that runs depends on the worker being fully ready.
+    time.sleep(8)
+
+    # Check if process is still running
+    if process.poll() is not None:
+        stderr_output = process.stderr.read() if process.stderr else ""
+        pytest.fail(f"Worker process exited prematurely. stderr: {stderr_output}")
+
+    # Provide worker info to tests
+    worker_info = {
+        "process": process,
+        "endpoint": dts_endpoint,
+        "taskhub": unique_taskhub,
+    }
+
+    try:
+        yield worker_info
+    finally:
+        # Cleanup: terminate worker subprocess
+        try:
+            process.terminate()
+            try:
+                process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait()
+        except Exception as e:
+            logging.warning(f"Error during worker process cleanup: {e}")
+
+
+@pytest.fixture(scope="module")
+def orchestration_helper(worker_process: dict[str, Any]) -> OrchestrationHelper:
+    """Create an OrchestrationHelper for the current test module."""
+    dts_client = create_dts_client(worker_process["endpoint"], worker_process["taskhub"])
+    return OrchestrationHelper(dts_client)
+
+
+@pytest.fixture(scope="module")
+def agent_client_factory(worker_process: dict[str, Any]) -> type[AgentClientFactoryProtocol]:
+    """Return a factory class for creating agent clients.
+
+    Usage in tests:
+        def test_example(self, agent_client_factory):
+            dts_client, agent_client = agent_client_factory.create(max_poll_retries=90)
+    """
+
+    class AgentClientFactory:
+        """Factory for creating DTS and Agent client pairs."""
+
+        endpoint = worker_process["endpoint"]
+        taskhub = worker_process["taskhub"]
+
+        @classmethod
+        def create(cls, max_poll_retries: int = 90) -> tuple[DurableTaskSchedulerClient, DurableAIAgentClient]:
+            """Create a DTS client and Agent client pair."""
+            return create_agent_client(cls.endpoint, cls.taskhub, max_poll_retries)
+
+    return AgentClientFactory
+
+
+@pytest.fixture(scope="module")
+def workflow_client(worker_process: dict[str, Any]) -> DurableWorkflowClient:
+    """Create a DurableWorkflowClient bound to the current sample worker's task hub."""
+    dts_client = create_dts_client(worker_process["endpoint"], worker_process["taskhub"])
+    return DurableWorkflowClient(dts_client)

python/packages/durabletask/tests/integration_tests/test_01_dt_single_agent.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for single agent functionality.
+
+Tests basic agent operations including:
+- Agent registration and retrieval
+- Single agent interactions
+- Conversation continuity across multiple messages
+- Multi-threaded agent usage
+- Empty thread ID handling
+"""
+
+from typing import Any, Protocol
+
+import pytest
+
+from agent_framework_durabletask import DurableAIAgentClient
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Module-level markers - applied to all tests in this module
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("01_single_agent"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_azure_openai,
+    pytest.mark.requires_dts,
+]
+
+
+class TestSingleAgent:
+    """Test suite for single agent functionality."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol]) -> None:
+        """Setup test fixtures."""
+        # Create agent client using the factory fixture
+        _, self.agent_client = agent_client_factory.create()
+
+    def test_agent_registration(self) -> None:
+        """Test that the Joker agent is registered and accessible."""
+        agent = self.agent_client.get_agent("Joker")
+        assert agent is not None
+        assert agent.name == "Joker"
+
+    def test_single_interaction(self):
+        """Test a single interaction with the agent."""
+        agent = self.agent_client.get_agent("Joker")
+        session = agent.create_session()
+
+        response = agent.run("Tell me a short joke about programming.", session=session)
+
+        assert response is not None
+        assert response.text is not None
+        assert len(response.text) > 0
+
+    def test_conversation_continuity(self):
+        """Test that conversation context is maintained across turns."""
+        agent = self.agent_client.get_agent("Joker")
+        session = agent.create_session()
+
+        # First turn: Ask for a joke about a specific topic
+        response1 = agent.run("Tell me a joke about cats.", session=session)
+        assert response1 is not None
+        assert len(response1.text) > 0
+
+        # Second turn: Ask a follow-up that requires context
+        response2 = agent.run("Can you make it funnier?", session=session)
+        assert response2 is not None
+        assert len(response2.text) > 0
+
+        # The agent should understand "it" refers to the previous joke
+
+    def test_multiple_sessions(self):
+        """Test that different sessions maintain separate contexts."""
+        agent = self.agent_client.get_agent("Joker")
+
+        # Create two separate sessions
+        session1 = agent.create_session()
+        session2 = agent.create_session()
+
+        assert session1.durable_session_id != session2.durable_session_id
+
+        # Send different messages to each session
+        response1 = agent.run("Tell me a joke about dogs.", session=session1)
+        response2 = agent.run("Tell me a joke about birds.", session=session2)
+
+        assert response1 is not None
+        assert response2 is not None
+        assert response1.text != response2.text

python/packages/durabletask/tests/integration_tests/test_02_dt_multi_agent.py

@@ -0,0 +1,113 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for multi-agent functionality.
+
+Tests operations with multiple specialized agents:
+- Multiple agent registration
+- Agent-specific tool usage
+- Independent thread management per agent
+- Concurrent agent operations
+- Agent isolation and tool routing
+"""
+
+from typing import Any, Protocol
+
+import pytest
+
+from agent_framework_durabletask import DurableAIAgentClient
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Agent names from the 02_multi_agent sample
+WEATHER_AGENT_NAME: str = "WeatherAgent"
+MATH_AGENT_NAME: str = "MathAgent"
+
+# Module-level markers - applied to all tests in this module
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("02_multi_agent"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_azure_openai,
+    pytest.mark.requires_dts,
+]
+
+
+class TestMultiAgent:
+    """Test suite for multi-agent functionality."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol]) -> None:
+        """Setup test fixtures."""
+        # Create agent client using the factory fixture
+        _, self.agent_client = agent_client_factory.create()
+
+    def test_multiple_agents_registered(self) -> None:
+        """Test that both agents are registered and accessible."""
+        weather_agent = self.agent_client.get_agent(WEATHER_AGENT_NAME)
+        math_agent = self.agent_client.get_agent(MATH_AGENT_NAME)
+
+        assert weather_agent is not None
+        assert weather_agent.name == WEATHER_AGENT_NAME
+        assert math_agent is not None
+        assert math_agent.name == MATH_AGENT_NAME
+
+    @pytest.mark.skip(reason="Flaky in CI: times out / crashes the xdist runner; temporarily disabled.")
+    def test_weather_agent_with_tool(self):
+        """Test weather agent with weather tool execution."""
+        agent = self.agent_client.get_agent(WEATHER_AGENT_NAME)
+        session = agent.create_session()
+
+        response = agent.run("What's the weather in Seattle?", session=session)
+
+        assert response is not None
+        assert response.text is not None
+        # Should contain weather information from the tool
+        assert len(response.text) > 0
+
+        # Verify that the get_weather tool was actually invoked
+        tool_calls = [
+            content for msg in response.messages for content in msg.contents if content.type == "function_call"
+        ]
+        assert len(tool_calls) > 0, "Expected at least one tool call"
+        assert any(call.name == "get_weather" for call in tool_calls), "Expected get_weather tool to be called"
+
+    @pytest.mark.skip(reason="Flaky in CI: times out / crashes the xdist runner; temporarily disabled.")
+    def test_math_agent_with_tool(self):
+        """Test math agent with calculation tool execution."""
+        agent = self.agent_client.get_agent(MATH_AGENT_NAME)
+        session = agent.create_session()
+
+        response = agent.run("Calculate a 20% tip on a $50 bill.", session=session)
+
+        assert response is not None
+        assert response.text is not None
+        # Should contain calculation results from the tool
+        assert len(response.text) > 0
+
+        # Verify that the calculate_tip tool was actually invoked
+        tool_calls = [
+            content for msg in response.messages for content in msg.contents if content.type == "function_call"
+        ]
+        assert len(tool_calls) > 0, "Expected at least one tool call"
+        assert any(call.name == "calculate_tip" for call in tool_calls), "Expected calculate_tip tool to be called"
+
+    def test_multiple_calls_to_same_agent(self):
+        """Test multiple sequential calls to the same agent."""
+        agent = self.agent_client.get_agent(WEATHER_AGENT_NAME)
+        session = agent.create_session()
+
+        # Multiple weather queries
+        response1 = agent.run("What's the weather in Chicago?", session=session)
+        response2 = agent.run("And what about Los Angeles?", session=session)
+
+        assert response1 is not None
+        assert response2 is not None
+        assert len(response1.text) > 0
+        assert len(response2.text) > 0

python/packages/durabletask/tests/integration_tests/test_03_dt_single_agent_streaming.py

@@ -0,0 +1,236 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Integration Tests for Reliable Streaming Sample
+
+Tests the reliable streaming sample using Redis Streams for persistent message delivery.
+
+The worker process is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/durabletask/tests/integration_tests/.env.example)
+- DTS emulator running (docker run -d -p 8080:8080 mcr.microsoft.com/durabletask/emulator:latest)
+- Redis running (docker run -d --name redis -p 6379:6379 redis:latest)
+
+Usage:
+    uv run pytest packages/durabletask/tests/integration_tests/test_03_single_agent_streaming.py -v
+"""
+
+import asyncio
+import os
+import sys
+import time
+from datetime import timedelta
+from pathlib import Path
+from typing import Any, Protocol
+
+import pytest
+import redis.asyncio as aioredis
+
+from agent_framework_durabletask import DurableAIAgentClient
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Add sample directory to path to import RedisStreamResponseHandler
+SAMPLE_DIR = Path(__file__).parents[4] / "samples" / "04-hosting" / "durabletask" / "03_single_agent_streaming"
+sys.path.insert(0, str(SAMPLE_DIR))
+
+from redis_stream_response_handler import (  # type: ignore[reportMissingImports] # pyrefly: ignore[missing-import] # ty: ignore[unresolved-import] # noqa: E402
+    RedisStreamResponseHandler,
+)
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("03_single_agent_streaming"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_azure_openai,
+    pytest.mark.requires_dts,
+    pytest.mark.requires_redis,
+]
+
+
+class TestSampleReliableStreaming:
+    """Tests for 03_single_agent_streaming sample."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol], orchestration_helper) -> None:
+        """Setup test fixtures."""
+        # Create agent client using the factory fixture
+        _, self.agent_client = agent_client_factory.create()
+        self.helper = orchestration_helper
+
+        # Redis configuration
+        self.redis_connection_string = os.environ.get("REDIS_CONNECTION_STRING", "redis://localhost:6379")
+        self.redis_stream_ttl_minutes = int(os.environ.get("REDIS_STREAM_TTL_MINUTES", "10"))
+
+    async def _get_stream_handler(self) -> RedisStreamResponseHandler:  # type: ignore[reportMissingTypeStubs]
+        """Create a new Redis stream handler for each request."""
+        redis_client = aioredis.from_url(  # type: ignore[reportUnknownMemberType]
+            self.redis_connection_string,
+            encoding="utf-8",
+            decode_responses=False,
+        )
+        return RedisStreamResponseHandler(  # type: ignore[reportUnknownMemberType]
+            redis_client=redis_client,
+            stream_ttl=timedelta(minutes=self.redis_stream_ttl_minutes),
+        )
+
+    async def _stream_from_redis(
+        self,
+        session_key: str,
+        cursor: str | None = None,
+        timeout: float = 30.0,
+    ) -> tuple[str, bool, str]:
+        """
+        Stream responses from Redis using the sample's RedisStreamResponseHandler.
+
+        Args:
+            session_key: The conversation/thread ID to stream from
+            cursor: Optional cursor to resume from
+            timeout: Maximum time to wait for stream completion
+
+        Returns:
+            Tuple of (accumulated text, completion status, last entry_id)
+        """
+        accumulated_text = ""
+        is_complete = False
+        last_entry_id = cursor if cursor else "0-0"
+        start_time = time.time()
+
+        async with await self._get_stream_handler() as stream_handler:  # type: ignore[reportUnknownMemberType]
+            try:
+                async for chunk in stream_handler.read_stream(session_key, cursor):  # type: ignore[reportUnknownMemberType]
+                    if time.time() - start_time > timeout:
+                        break
+
+                    last_entry_id = chunk.entry_id  # type: ignore[reportUnknownMemberType]
+
+                    if chunk.error:  # type: ignore[reportUnknownMemberType]
+                        # Stream not found or timeout - this is expected if agent hasn't written yet
+                        # Don't raise an error, just return what we have
+                        break
+
+                    if chunk.is_done:  # type: ignore[reportUnknownMemberType]
+                        is_complete = True
+                        break
+
+                    if chunk.text:  # type: ignore[reportUnknownMemberType]
+                        accumulated_text += chunk.text  # type: ignore[reportUnknownMemberType]
+
+            except Exception as ex:
+                # For test purposes, we catch exceptions and return what we have
+                if "timed out" not in str(ex).lower():
+                    raise
+
+        return accumulated_text, is_complete, last_entry_id  # type: ignore[reportReturnType]
+
+    def test_agent_run_and_stream(self) -> None:
+        """Test agent execution with Redis streaming."""
+        # Get the TravelPlanner agent
+        travel_planner = self.agent_client.get_agent("TravelPlanner")
+        assert travel_planner is not None
+        assert travel_planner.name == "TravelPlanner"
+
+        # Create a new session
+        session = travel_planner.create_session()
+        assert session.durable_session_id is not None
+        assert session.durable_session_id.key is not None
+        session_key = str(session.durable_session_id.key)
+
+        # Start agent run with wait_for_response=False for non-blocking execution
+        travel_planner.run(
+            "Plan a 1-day trip to Seattle in 1 sentence", session=session, options={"wait_for_response": False}
+        )
+
+        # Poll Redis stream with retries to handle race conditions
+        # The agent may take a few seconds to process and start writing to Redis
+        # We use cursor-based resumption to continue reading from where we left off
+        max_retries = 20
+        retry_count = 0
+        accumulated_text = ""
+        is_complete = False
+        cursor: str | None = None
+
+        while retry_count < max_retries and not is_complete:
+            text, is_complete, last_cursor = asyncio.run(
+                self._stream_from_redis(session_key, cursor=cursor, timeout=10.0)
+            )
+            accumulated_text += text
+            cursor = last_cursor  # Resume from last position on next read
+
+            if is_complete:
+                # Stream completed successfully
+                break
+
+            if len(accumulated_text) > 0:
+                # Got content but not completion marker yet - keep reading without delay
+                # The agent may still be streaming or about to write completion marker
+                continue
+
+            # No content yet - wait before retrying
+            time.sleep(2)
+            retry_count += 1
+
+        # Verify we got content
+        assert len(accumulated_text) > 0, (
+            f"Expected text content but got empty string for session_key: {session_key} after {retry_count} retries"
+        )
+        assert "seattle" in accumulated_text.lower(), f"Expected 'seattle' in response but got: {accumulated_text}"
+        assert is_complete, "Expected stream to be complete"
+
+    def test_stream_with_cursor_resumption(self) -> None:
+        """Test streaming with cursor-based resumption."""
+        # Get the TravelPlanner agent
+        travel_planner = self.agent_client.get_agent("TravelPlanner")
+        session = travel_planner.create_session()
+        assert session.durable_session_id is not None
+        assert session.durable_session_id.key is not None
+        session_key = str(session.durable_session_id.key)
+
+        # Start agent run
+        travel_planner.run("What's the weather like?", session=session, options={"wait_for_response": False})
+
+        # Wait for agent to start writing
+        time.sleep(3)
+
+        # Read partial stream to get a cursor
+        async def get_partial_stream() -> tuple[str, str]:
+            async with await self._get_stream_handler() as stream_handler:  # type: ignore[reportUnknownMemberType]
+                accumulated_text = ""
+                last_entry_id = "0-0"
+                chunk_count = 0
+
+                # Read just first 2 chunks
+                async for chunk in stream_handler.read_stream(session_key):  # type: ignore[reportUnknownMemberType]
+                    last_entry_id = chunk.entry_id  # type: ignore[reportUnknownMemberType]
+                    if chunk.text:  # type: ignore[reportUnknownMemberType]
+                        accumulated_text += chunk.text  # type: ignore[reportUnknownMemberType]
+                    chunk_count += 1
+                    if chunk_count >= 2:
+                        break
+
+                return accumulated_text, last_entry_id  # type: ignore[reportReturnType]
+
+        partial_text, cursor = asyncio.run(get_partial_stream())
+
+        # Resume from cursor
+        remaining_text, _, _ = asyncio.run(self._stream_from_redis(session_key, cursor=cursor))
+
+        # Verify we got some initial content
+        assert len(partial_text) > 0
+
+        # Combined text should be coherent
+        full_text = partial_text + remaining_text
+        assert len(full_text) > 0
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/packages/durabletask/tests/integration_tests/test_04_dt_single_agent_orchestration_chaining.py

@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for single agent orchestration with chaining.
+
+Tests orchestration patterns with sequential agent calls:
+- Orchestration registration and execution
+- Sequential agent calls on same thread
+- Conversation continuity in orchestrations
+- Thread context preservation
+"""
+
+import json
+import logging
+from typing import Any, Protocol
+
+import pytest
+from durabletask.client import OrchestrationStatus
+
+from agent_framework_durabletask import DurableAIAgentClient
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Agent name from the 04_single_agent_orchestration_chaining sample
+WRITER_AGENT_NAME: str = "WriterAgent"
+
+# Configure logging
+logging.basicConfig(level=logging.WARNING)
+
+# Module-level markers - applied to all tests in this module
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("04_single_agent_orchestration_chaining"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_azure_openai,
+    pytest.mark.requires_dts,
+]
+
+
+class TestSingleAgentOrchestrationChaining:
+    """Test suite for single agent orchestration with chaining."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol], orchestration_helper) -> None:
+        """Setup test fixtures."""
+        # Create agent client using the factory fixture
+        self.dts_client, self.agent_client = agent_client_factory.create()
+        self.orch_helper = orchestration_helper
+
+    def test_agent_registered(self):
+        """Test that the Writer agent is registered."""
+        agent = self.agent_client.get_agent(WRITER_AGENT_NAME)
+        assert agent is not None
+        assert agent.name == WRITER_AGENT_NAME
+
+    def test_chaining_context_preserved(self):
+        """Test that context is preserved across agent runs in orchestration."""
+        # Start the orchestration
+        instance_id = self.dts_client.schedule_new_orchestration(
+            orchestrator="single_agent_chaining_orchestration",
+            input="",
+        )
+
+        # Wait for completion with output
+        metadata, output = self.orch_helper.wait_for_orchestration_with_output(
+            instance_id=instance_id,
+            timeout=120.0,
+        )
+
+        assert metadata is not None
+        assert output is not None
+
+        # The final output should be a refined sentence
+        final_text = json.loads(output)
+
+        # Should be a meaningful sentence (not empty or error message)
+        assert len(final_text) > 10
+        assert not final_text.startswith("Error")
+
+    def test_multiple_orchestration_instances(self):
+        """Test that multiple orchestration instances can run independently."""
+        # Start two orchestrations
+        instance_id_1 = self.dts_client.schedule_new_orchestration(
+            orchestrator="single_agent_chaining_orchestration",
+            input="",
+        )
+        instance_id_2 = self.dts_client.schedule_new_orchestration(
+            orchestrator="single_agent_chaining_orchestration",
+            input="",
+        )
+
+        assert instance_id_1 != instance_id_2
+
+        # Both should complete
+        metadata_1 = self.orch_helper.wait_for_orchestration(
+            instance_id=instance_id_1,
+            timeout=120.0,
+        )
+        metadata_2 = self.orch_helper.wait_for_orchestration(
+            instance_id=instance_id_2,
+            timeout=120.0,
+        )
+
+        assert metadata_1.runtime_status == OrchestrationStatus.COMPLETED
+        assert metadata_2.runtime_status == OrchestrationStatus.COMPLETED

python/packages/durabletask/tests/integration_tests/test_05_dt_multi_agent_orchestration_concurrency.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for multi-agent orchestration with concurrency.
+
+Tests concurrent execution patterns:
+- Parallel agent execution
+- Concurrent orchestration tasks
+- Independent thread management in parallel
+- Result aggregation from concurrent calls
+"""
+
+import json
+import logging
+from typing import Any, Protocol
+
+import pytest
+from durabletask.client import OrchestrationStatus
+
+from agent_framework_durabletask import DurableAIAgentClient
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Agent names from the 05_multi_agent_orchestration_concurrency sample
+PHYSICIST_AGENT_NAME: str = "PhysicistAgent"
+CHEMIST_AGENT_NAME: str = "ChemistAgent"
+
+# Configure logging
+logging.basicConfig(level=logging.WARNING)
+
+# Module-level markers
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("05_multi_agent_orchestration_concurrency"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_dts,
+]
+
+
+class TestMultiAgentOrchestrationConcurrency:
+    """Test suite for multi-agent orchestration with concurrency."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol], orchestration_helper) -> None:
+        """Setup test fixtures."""
+        # Create agent client using the factory fixture
+        self.dts_client, self.agent_client = agent_client_factory.create()
+        self.orch_helper = orchestration_helper
+
+    def test_agents_registered(self):
+        """Test that both agents are registered."""
+        physicist = self.agent_client.get_agent(PHYSICIST_AGENT_NAME)
+        chemist = self.agent_client.get_agent(CHEMIST_AGENT_NAME)
+
+        assert physicist is not None
+        assert physicist.name == PHYSICIST_AGENT_NAME
+        assert chemist is not None
+        assert chemist.name == CHEMIST_AGENT_NAME
+
+    def test_different_prompts(self):
+        """Test concurrent orchestration with different prompts."""
+        prompts = [
+            "What is temperature?",
+            "Explain molecules.",
+        ]
+
+        for prompt in prompts:
+            instance_id = self.dts_client.schedule_new_orchestration(
+                orchestrator="multi_agent_concurrent_orchestration",
+                input=prompt,
+            )
+
+            metadata, output = self.orch_helper.wait_for_orchestration_with_output(
+                instance_id=instance_id,
+                timeout=120.0,
+            )
+
+            assert metadata.runtime_status == OrchestrationStatus.COMPLETED
+            result = json.loads(output)
+            assert "physicist" in result
+            assert "chemist" in result

python/packages/durabletask/tests/integration_tests/test_06_dt_multi_agent_orchestration_conditionals.py

@@ -0,0 +1,86 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for multi-agent orchestration with conditionals.
+
+Tests conditional orchestration patterns:
+- Conditional branching in orchestrations
+- Agent-based decision making
+- Activity function execution
+- Structured output handling
+- Conditional routing based on agent responses
+"""
+
+import logging
+from typing import Any, Protocol
+
+import pytest
+from durabletask.client import OrchestrationStatus
+
+from agent_framework_durabletask import DurableAIAgentClient
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Agent names from the 06_multi_agent_orchestration_conditionals sample
+SPAM_AGENT_NAME: str = "SpamDetectionAgent"
+EMAIL_AGENT_NAME: str = "EmailAssistantAgent"
+
+# Configure logging
+logging.basicConfig(level=logging.WARNING)
+
+# Module-level markers
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("06_multi_agent_orchestration_conditionals"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_dts,
+]
+
+
+class TestMultiAgentOrchestrationConditionals:
+    """Test suite for multi-agent orchestration with conditionals."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol], orchestration_helper) -> None:
+        """Setup test fixtures."""
+        # Create agent client using the factory fixture
+        self.dts_client, self.agent_client = agent_client_factory.create()
+        self.orch_helper = orchestration_helper
+
+    def test_agents_registered(self):
+        """Test that both agents are registered."""
+        spam_agent = self.agent_client.get_agent(SPAM_AGENT_NAME)
+        email_agent = self.agent_client.get_agent(EMAIL_AGENT_NAME)
+
+        assert spam_agent is not None
+        assert spam_agent.name == SPAM_AGENT_NAME
+        assert email_agent is not None
+        assert email_agent.name == EMAIL_AGENT_NAME
+
+    @pytest.mark.skip(reason="Flaky in CI: times out / crashes the xdist runner; temporarily disabled.")
+    def test_conditional_branching(self):
+        """Test that conditional branching works correctly."""
+        # Test with obvious spam
+        spam_payload = {
+            "email_id": "spam-001",
+            "email_content": "Buy cheap medications online! No prescription needed! Limited time offer!",
+        }
+
+        spam_instance_id = self.dts_client.schedule_new_orchestration(
+            orchestrator="spam_detection_orchestration",
+            input=spam_payload,
+        )
+
+        # Both should complete successfully (different branches)
+        spam_metadata = self.orch_helper.wait_for_orchestration(
+            instance_id=spam_instance_id,
+            timeout=120.0,
+        )
+
+        assert spam_metadata.runtime_status == OrchestrationStatus.COMPLETED

python/packages/durabletask/tests/integration_tests/test_07_dt_single_agent_orchestration_hitl.py

@@ -0,0 +1,174 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for single agent orchestration with human-in-the-loop.
+
+Tests human-in-the-loop (HITL) patterns:
+- External event waiting and handling
+- Timeout handling in orchestrations
+- Iterative refinement with human feedback
+- Activity function integration
+- Approval workflow patterns
+"""
+
+import logging
+from typing import Any, Protocol
+
+import pytest
+from durabletask.client import OrchestrationStatus
+
+from agent_framework_durabletask import DurableAIAgentClient
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Constants from the 07_single_agent_orchestration_hitl sample
+WRITER_AGENT_NAME: str = "WriterAgent"
+HUMAN_APPROVAL_EVENT: str = "HumanApproval"
+
+# Configure logging
+logging.basicConfig(level=logging.WARNING)
+
+# Module-level markers
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("07_single_agent_orchestration_hitl"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_dts,
+]
+
+
+class TestSingleAgentOrchestrationHITL:
+    """Test suite for single agent orchestration with human-in-the-loop."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol], orchestration_helper) -> None:
+        """Setup test fixtures."""
+        # Create agent client using the factory fixture
+        self.dts_client, self.agent_client = agent_client_factory.create()
+        self.orch_helper = orchestration_helper
+
+    def test_agent_registered(self):
+        """Test that the Writer agent is registered."""
+        agent = self.agent_client.get_agent(WRITER_AGENT_NAME)
+        assert agent is not None
+        assert agent.name == WRITER_AGENT_NAME
+
+    def test_hitl_orchestration_with_approval(self):
+        """Test HITL orchestration with immediate approval."""
+        payload = {
+            "topic": "The benefits of continuous learning",
+            "max_review_attempts": 3,
+            "approval_timeout_seconds": 60,
+        }
+
+        # Start the orchestration
+        instance_id = self.dts_client.schedule_new_orchestration(
+            orchestrator="content_generation_hitl_orchestration",
+            input=payload,
+        )
+
+        assert instance_id is not None
+
+        # Wait for orchestration to reach notification point
+        notification_received = self.orch_helper.wait_for_notification(instance_id, timeout_seconds=90)
+        assert notification_received, "Failed to receive notification from orchestration"
+
+        # Send approval event
+        approval_data = {"approved": True, "feedback": ""}
+        self.orch_helper.raise_event(
+            instance_id=instance_id,
+            event_name=HUMAN_APPROVAL_EVENT,
+            event_data=approval_data,
+        )
+
+        # Wait for completion
+        metadata = self.orch_helper.wait_for_orchestration(
+            instance_id=instance_id,
+            timeout=90.0,
+        )
+
+        assert metadata is not None
+        assert metadata.runtime_status == OrchestrationStatus.COMPLETED
+
+    def test_hitl_orchestration_with_rejection_and_feedback(self):
+        """Test HITL orchestration with rejection and iterative refinement."""
+        payload = {
+            "topic": "Artificial Intelligence in healthcare",
+            "max_review_attempts": 3,
+            "approval_timeout_seconds": 60,
+        }
+
+        # Start the orchestration
+        instance_id = self.dts_client.schedule_new_orchestration(
+            orchestrator="content_generation_hitl_orchestration",
+            input=payload,
+        )
+
+        # Wait for orchestration to reach notification point
+        notification_received = self.orch_helper.wait_for_notification(instance_id, timeout_seconds=90)
+        assert notification_received, "Failed to receive notification from orchestration"
+
+        # First rejection with feedback
+        rejection_data = {
+            "approved": False,
+            "feedback": "Please make it more concise and add specific examples.",
+        }
+        self.orch_helper.raise_event(
+            instance_id=instance_id,
+            event_name=HUMAN_APPROVAL_EVENT,
+            event_data=rejection_data,
+        )
+
+        # Wait for orchestration to refine and reach notification point again
+        notification_received = self.orch_helper.wait_for_notification(instance_id, timeout_seconds=90)
+        assert notification_received, "Failed to receive notification after refinement"
+
+        # Second approval
+        approval_data = {"approved": True, "feedback": ""}
+        self.orch_helper.raise_event(
+            instance_id=instance_id,
+            event_name=HUMAN_APPROVAL_EVENT,
+            event_data=approval_data,
+        )
+
+        # Wait for completion
+        metadata = self.orch_helper.wait_for_orchestration(
+            instance_id=instance_id,
+            timeout=90.0,
+        )
+
+        assert metadata is not None
+        assert metadata.runtime_status == OrchestrationStatus.COMPLETED
+
+    def test_hitl_orchestration_timeout(self):
+        """Test HITL orchestration timeout behavior."""
+        payload = {
+            "topic": "Cloud computing fundamentals",
+            "max_review_attempts": 1,
+            "approval_timeout_seconds": 0.1,  # Short timeout for testing
+        }
+
+        # Start the orchestration
+        instance_id = self.dts_client.schedule_new_orchestration(
+            orchestrator="content_generation_hitl_orchestration",
+            input=payload,
+        )
+
+        # Don't send any approval - let it timeout
+        # The orchestration should fail due to timeout
+        try:
+            metadata = self.orch_helper.wait_for_orchestration(
+                instance_id=instance_id,
+                timeout=90.0,
+            )
+            # If it completes, it should be failed status due to timeout
+            assert metadata.runtime_status == OrchestrationStatus.FAILED
+        except (RuntimeError, TimeoutError):
+            # Expected - orchestration should timeout and fail
+            pass

python/packages/durabletask/tests/integration_tests/test_08_dt_workflow.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for the standalone durabletask workflow sample (08_workflow).
+
+Exercises the standalone (non-Azure-Functions) workflow path:
+- ``DurableAIAgentWorker.configure_workflow`` auto-registers the agent entities,
+  non-agent executor activities, and the workflow orchestrator.
+- A client starts the workflow by scheduling its ``dafx-{workflow_name}`` orchestration.
+- Conditional routing sends spam to a non-agent handler and legitimate email
+  through a second agent and a sender executor.
+"""
+
+import logging
+from typing import Any, Protocol
+
+import pytest
+from durabletask.client import OrchestrationStatus
+
+from agent_framework_durabletask import DurableAIAgentClient, workflow_orchestrator_name
+
+# Must match the workflow name in samples/04-hosting/durabletask/08_workflow/worker.py
+WORKFLOW_NAME = "email_triage"
+
+logging.basicConfig(level=logging.WARNING)
+
+
+class AgentClientFactoryProtocol(Protocol):
+    """Protocol for the agent client factory fixture."""
+
+    @classmethod
+    def create(cls, max_poll_retries: int = 90) -> tuple[Any, DurableAIAgentClient]: ...
+
+
+# Module-level markers
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("08_workflow"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_dts,
+]
+
+
+class TestStandaloneWorkflow:
+    """Standalone (non-Azure-Functions) workflow execution on a durabletask worker."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, agent_client_factory: type[AgentClientFactoryProtocol], orchestration_helper) -> None:
+        """Provide a DTS client and orchestration helper for each test."""
+        self.dts_client, self.agent_client = agent_client_factory.create()
+        self.orch_helper = orchestration_helper
+
+    def test_legitimate_email_drafts_response(self) -> None:
+        """A legitimate email routes through the email agent and is 'sent'."""
+        instance_id = self.dts_client.schedule_new_orchestration(
+            orchestrator=workflow_orchestrator_name(WORKFLOW_NAME),
+            input=(
+                "Hi team, just a reminder about our sprint planning meeting tomorrow at 10 AM. "
+                "Please review the agenda in Jira."
+            ),
+        )
+
+        metadata, output = self.orch_helper.wait_for_orchestration_with_output(
+            instance_id=instance_id,
+            timeout=180.0,
+        )
+
+        assert metadata.runtime_status == OrchestrationStatus.COMPLETED
+        assert output is not None
+        assert "Email sent" in str(output)
+
+    def test_spam_email_handled(self) -> None:
+        """A spam email routes to the non-agent spam handler."""
+        instance_id = self.dts_client.schedule_new_orchestration(
+            orchestrator=workflow_orchestrator_name(WORKFLOW_NAME),
+            input="URGENT! You've won $1,000,000! Click here now to claim your prize! Limited time offer!",
+        )
+
+        metadata, output = self.orch_helper.wait_for_orchestration_with_output(
+            instance_id=instance_id,
+            timeout=180.0,
+        )
+
+        assert metadata.runtime_status == OrchestrationStatus.COMPLETED
+        assert output is not None
+        assert "spam" in str(output).lower()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/packages/durabletask/tests/integration_tests/test_09_dt_workflow_hitl.py

@@ -0,0 +1,112 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for the standalone durabletask HITL workflow sample (09_workflow_hitl).
+
+Exercises the human-in-the-loop workflow path on a standalone durabletask worker:
+- The ``InputRouter`` start executor receives a typed ``ContentSubmission`` that the
+  shared engine reconstructs from the client's JSON payload (no manual parsing).
+- An analysis agent produces a recommendation, then the workflow pauses for human
+  approval via ``request_info``.
+- The client retrieves the pending request, replies with ``send_hitl_response``, and
+  the workflow resumes to an approved/rejected outcome read via ``await_workflow_output``.
+"""
+
+import logging
+import time
+from typing import Any
+
+import pytest
+
+from agent_framework_durabletask import DurableWorkflowClient
+
+logging.basicConfig(level=logging.WARNING)
+
+# Must match the workflow name in samples/04-hosting/durabletask/09_workflow_hitl/worker.py
+WORKFLOW_NAME = "content_moderation"
+
+# Module-level markers
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("09_workflow_hitl"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_dts,
+    pytest.mark.requires_azure_openai,
+]
+
+
+def _wait_for_hitl_request(
+    client: DurableWorkflowClient, instance_id: str, timeout_seconds: int = 90
+) -> list[dict[str, Any]]:
+    """Poll until the workflow records at least one pending HITL request."""
+    deadline = time.time() + timeout_seconds
+    while time.time() < deadline:
+        pending = client.get_pending_hitl_requests(instance_id, workflow_name=WORKFLOW_NAME)
+        if pending:
+            return pending
+        time.sleep(2)
+    raise AssertionError(f"Timed out waiting for a HITL request on instance {instance_id}")
+
+
+class TestStandaloneWorkflowHITL:
+    """Human-in-the-loop workflow execution on a standalone durabletask worker."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, workflow_client: DurableWorkflowClient) -> None:
+        """Bind the DurableWorkflowClient for the current sample worker."""
+        self.client = workflow_client
+
+    def _run_case(self, submission: dict[str, Any], *, approve: bool) -> Any:
+        """Start a moderation case, answer the HITL pause, and return the final output."""
+        instance_id = self.client.start_workflow(input=submission, workflow_name=WORKFLOW_NAME)
+
+        pending = _wait_for_hitl_request(self.client, instance_id)
+        request = pending[0]
+        assert request["request_id"]
+        assert request["source_executor_id"]
+
+        self.client.send_hitl_response(
+            instance_id,
+            request["request_id"],
+            {"approved": approve, "reviewer_notes": "Looks good." if approve else "Violates content policy."},
+            workflow_name=WORKFLOW_NAME,
+        )
+
+        return self.client.await_workflow_output(instance_id, workflow_name=WORKFLOW_NAME, timeout_seconds=180)
+
+    def test_hitl_workflow_approval(self) -> None:
+        """Appropriate content is approved after the reviewer says yes."""
+        output = self._run_case(
+            {
+                "content_id": "article-001",
+                "title": "Introduction to AI in Healthcare",
+                "body": (
+                    "Artificial intelligence is improving healthcare by enabling faster diagnosis, "
+                    "personalized treatment plans, and better patient outcomes."
+                ),
+                "author": "Dr. Jane Smith",
+            },
+            approve=True,
+        )
+
+        assert output is not None
+        assert "APPROVED" in str(output).upper()
+
+    def test_hitl_workflow_rejection(self) -> None:
+        """Spammy content is rejected after the reviewer says no."""
+        output = self._run_case(
+            {
+                "content_id": "article-002",
+                "title": "Get Rich Quick",
+                "body": "Click here NOW to make $10,000 overnight! GUARANTEED! Limited time offer!",
+                "author": "Definitely Not Spam",
+            },
+            approve=False,
+        )
+
+        assert output is not None
+        assert "REJECTED" in str(output).upper()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/packages/durabletask/tests/integration_tests/test_11_dt_subworkflow.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for the composed sub-workflow sample (11_subworkflow).
+
+Exercises workflow *composition* on a standalone durabletask worker:
+- An outer ``review_pipeline`` embeds an inner ``sentiment_analysis`` workflow via a
+  ``WorkflowExecutor`` node (``sentiment_sub``).
+- ``DurableAIAgentWorker.configure_workflow`` walks the composition and registers a
+  durable orchestration for each workflow; the inner workflow runs as a child
+  orchestration when the outer reaches the ``WorkflowExecutor`` node.
+- The inner workflow's output (a sentiment summary) is forwarded to the outer
+  ``reporter`` executor, which produces the final result.
+
+The inner workflow hosts an AI agent, so these tests require model credentials.
+"""
+
+import logging
+from typing import Any
+
+import pytest
+
+from agent_framework_durabletask import DurableWorkflowClient
+
+logging.basicConfig(level=logging.WARNING)
+
+# Must match the outer workflow name in samples/04-hosting/durabletask/11_subworkflow/worker.py
+WORKFLOW_NAME = "review_pipeline"
+
+# Module-level markers
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("11_subworkflow"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_dts,
+    pytest.mark.requires_azure_openai,
+]
+
+
+class TestSubworkflowComposition:
+    """Composed (outer + inner) workflow execution on a standalone durabletask worker."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, workflow_client: DurableWorkflowClient) -> None:
+        """Bind the DurableWorkflowClient for the current sample worker."""
+        self.client = workflow_client
+
+    def _run(self, review: str) -> Any:
+        """Run the composed workflow with a review and return its final output."""
+        instance_id = self.client.start_workflow(input=review, workflow_name=WORKFLOW_NAME)
+        return self.client.await_workflow_output(instance_id, workflow_name=WORKFLOW_NAME, timeout_seconds=180)
+
+    def test_positive_review_runs_through_subworkflow(self) -> None:
+        """A positive review flows through the embedded sentiment sub-workflow to a report."""
+        output = self._run(
+            "Absolutely love this espresso machine - it heats up fast and the coffee is consistently great."
+        )
+
+        assert output is not None
+        # The outer reporter wraps the inner sub-workflow's forwarded sentiment summary.
+        assert "sentiment" in str(output).lower()
+
+    def test_negative_review_runs_through_subworkflow(self) -> None:
+        """A negative review also completes the composed pipeline end-to-end."""
+        output = self._run("Disappointed. The device stopped working after two weeks and support never replied.")
+
+        assert output is not None
+        assert "sentiment" in str(output).lower()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/packages/durabletask/tests/integration_tests/test_12_dt_subworkflow_hitl.py

@@ -0,0 +1,152 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Integration tests for the composed sub-workflow HITL sample (12_subworkflow_hitl).
+
+Exercises human-in-the-loop **inside a nested sub-workflow** on a standalone
+durabletask worker:
+- An outer ``moderation_pipeline`` embeds an inner ``human_review`` workflow via a
+  ``WorkflowExecutor`` node (``review_sub``); on the durable host the inner workflow
+  runs as a child orchestration.
+- The inner ``review_gate`` pauses via ``request_info``. The pending request surfaces
+  at the top-level instance with a **qualified** id ``review_sub~0~{requestId}`` (the
+  ``~{ordinal}~`` hop addresses the specific child the node dispatched).
+- The client responds with that qualified id against the *top-level* instance and the
+  host routes it to the owning child orchestration, resuming to an approved/rejected
+  outcome.
+
+This sample hosts **no AI agents**, so it needs only the DTS emulator (no model
+credentials), which makes it a deterministic end-to-end check of the nested-HITL
+addressing.
+"""
+
+import logging
+import time
+from typing import Any
+
+import pytest
+
+from agent_framework_durabletask import DurableWorkflowClient
+from agent_framework_durabletask._workflows.naming import SUBWORKFLOW_REQUEST_SEPARATOR
+
+logging.basicConfig(level=logging.WARNING)
+
+# Must match the outer workflow name in samples/04-hosting/durabletask/12_subworkflow_hitl/worker.py
+WORKFLOW_NAME = "moderation_pipeline"
+# The WorkflowExecutor node id that embeds the inner HITL workflow.
+SUBWORKFLOW_NODE_ID = "review_sub"
+
+# Module-level markers. No requires_azure_openai: the sample hosts no agents.
+pytestmark = [
+    pytest.mark.flaky,
+    pytest.mark.integration,
+    pytest.mark.sample("12_subworkflow_hitl"),
+    pytest.mark.integration_test,
+    pytest.mark.requires_dts,
+]
+
+
+def _wait_for_hitl_request(
+    client: DurableWorkflowClient, instance_id: str, timeout_seconds: int = 90
+) -> list[dict[str, Any]]:
+    """Poll until the workflow (or a nested sub-workflow) records a pending HITL request."""
+    deadline = time.time() + timeout_seconds
+    while time.time() < deadline:
+        pending = client.get_pending_hitl_requests(instance_id, workflow_name=WORKFLOW_NAME)
+        if pending:
+            return pending
+        time.sleep(2)
+    raise AssertionError(f"Timed out waiting for a nested HITL request on instance {instance_id}")
+
+
+class TestSubworkflowHITL:
+    """Nested (sub-workflow) human-in-the-loop on a standalone durabletask worker."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, workflow_client: DurableWorkflowClient) -> None:
+        """Bind the DurableWorkflowClient for the current sample worker."""
+        self.client = workflow_client
+
+    def _run_case(self, submission: dict[str, Any], *, approve: bool) -> tuple[dict[str, Any], Any]:
+        """Start a moderation case, answer the nested HITL pause, return (request, output)."""
+        instance_id = self.client.start_workflow(input=submission, workflow_name=WORKFLOW_NAME)
+
+        pending = _wait_for_hitl_request(self.client, instance_id)
+        request = pending[0]
+
+        self.client.send_hitl_response(
+            instance_id,
+            request["request_id"],
+            {"approved": approve, "reviewer_notes": "Looks good." if approve else "Violates content policy."},
+            workflow_name=WORKFLOW_NAME,
+        )
+
+        output = self.client.await_workflow_output(instance_id, workflow_name=WORKFLOW_NAME, timeout_seconds=180)
+        return request, output
+
+    def test_nested_request_id_is_qualified_with_ordinal(self) -> None:
+        """The nested pending request surfaces with a ``review_sub~0~{id}`` qualified id."""
+        instance_id = self.client.start_workflow(
+            input={
+                "content_id": "article-100",
+                "title": "Quarterly Roadmap",
+                "body": "A summary of the upcoming features planned for the next quarter.",
+            },
+            workflow_name=WORKFLOW_NAME,
+        )
+
+        pending = _wait_for_hitl_request(self.client, instance_id)
+
+        assert len(pending) == 1
+        request = pending[0]
+        # The qualifier carries the node id and the child's ordinal (0 for the single
+        # dispatch), then the inner bare request id: ``review_sub~0~{requestId}``.
+        expected_prefix = f"{SUBWORKFLOW_NODE_ID}{SUBWORKFLOW_REQUEST_SEPARATOR}0{SUBWORKFLOW_REQUEST_SEPARATOR}"
+        assert request["request_id"].startswith(expected_prefix), request["request_id"]
+        # The bare inner id is non-empty after the qualifier.
+        assert request["request_id"][len(expected_prefix) :]
+        # The originating executor is the inner workflow's review gate.
+        assert request["source_executor_id"] == "review_gate"
+
+        # Drain the pause so the worker does not leave the instance hanging.
+        self.client.send_hitl_response(
+            instance_id,
+            request["request_id"],
+            {"approved": True, "reviewer_notes": "ok"},
+            workflow_name=WORKFLOW_NAME,
+        )
+        self.client.await_workflow_output(instance_id, workflow_name=WORKFLOW_NAME, timeout_seconds=180)
+
+    def test_nested_hitl_approval(self) -> None:
+        """Responding 'approved' to the nested request resumes the outer workflow to APPROVED."""
+        _request, output = self._run_case(
+            {
+                "content_id": "article-001",
+                "title": "Introduction to AI in Healthcare",
+                "body": (
+                    "Artificial intelligence is improving healthcare by enabling faster diagnosis, "
+                    "personalized treatment plans, and better patient outcomes."
+                ),
+            },
+            approve=True,
+        )
+
+        assert output is not None
+        assert "APPROVED" in str(output).upper()
+
+    def test_nested_hitl_rejection(self) -> None:
+        """Responding 'rejected' to the nested request resumes the outer workflow to REJECTED."""
+        _request, output = self._run_case(
+            {
+                "content_id": "article-002",
+                "title": "Get Rich Quick",
+                "body": "Click here NOW to make $10,000 overnight! GUARANTEED! Limited time offer!",
+            },
+            approve=False,
+        )
+
+        assert output is not None
+        assert "REJECTED" in str(output).upper()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/packages/durabletask/tests/test_agent_session_id.py

@@ -0,0 +1,310 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for AgentSessionId and DurableAgentSession."""
+
+from typing import Any
+
+import pytest
+from agent_framework import AgentSession
+
+from agent_framework_durabletask._models import AgentSessionId, DurableAgentSession
+
+
+class TestAgentSessionId:
+    """Test suite for AgentSessionId."""
+
+    def test_init_creates_session_id(self) -> None:
+        """Test that AgentSessionId initializes correctly."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key-123"
+
+    def test_with_random_key_generates_guid(self) -> None:
+        """Test that with_random_key generates a GUID."""
+        session_id = AgentSessionId.with_random_key(name="AgentEntity")
+
+        assert session_id.name == "AgentEntity"
+        assert len(session_id.key) == 32  # UUID hex is 32 chars
+        # Verify it's a valid hex string
+        int(session_id.key, 16)
+
+    def test_with_random_key_unique_keys(self) -> None:
+        """Test that with_random_key generates unique keys."""
+        session_id1 = AgentSessionId.with_random_key(name="AgentEntity")
+        session_id2 = AgentSessionId.with_random_key(name="AgentEntity")
+
+        assert session_id1.key != session_id2.key
+
+    def test_str_representation(self) -> None:
+        """Test string representation."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
+        str_repr = str(session_id)
+
+        assert str_repr == "@AgentEntity@test-key-123"
+
+    def test_repr_representation(self) -> None:
+        """Test repr representation."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key")
+        repr_str = repr(session_id)
+
+        assert "AgentSessionId" in repr_str
+        assert "AgentEntity" in repr_str
+        assert "test-key" in repr_str
+
+    def test_parse_valid_session_id(self) -> None:
+        """Test parsing valid session ID string."""
+        session_id = AgentSessionId.parse("@AgentEntity@test-key-123")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key-123"
+
+    def test_parse_invalid_format_no_prefix(self) -> None:
+        """Test parsing invalid format without @ prefix."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("AgentEntity@test-key")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
+    def test_parse_invalid_format_single_part(self) -> None:
+        """Test parsing invalid format with single part."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("@AgentEntity")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
+    def test_parse_with_multiple_at_signs_in_key(self) -> None:
+        """Test parsing with @ signs in the key."""
+        session_id = AgentSessionId.parse("@AgentEntity@key-with@symbols")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "key-with@symbols"
+
+    def test_parse_round_trip(self) -> None:
+        """Test round-trip parse and string conversion."""
+        original = AgentSessionId(name="AgentEntity", key="test-key")
+        str_repr = str(original)
+        parsed = AgentSessionId.parse(str_repr)
+
+        assert parsed.name == original.name
+        assert parsed.key == original.key
+
+    def test_to_entity_name_adds_prefix(self) -> None:
+        """Test that to_entity_name adds the dafx- prefix."""
+        entity_name = AgentSessionId.to_entity_name("TestAgent")
+        assert entity_name == "dafx-TestAgent"
+
+    def test_parse_with_agent_name_override(self) -> None:
+        """Test parsing @name@key format with agent_name parameter overrides the name."""
+        session_id = AgentSessionId.parse("@OriginalAgent@test-key-123", agent_name="OverriddenAgent")
+
+        assert session_id.name == "OverriddenAgent"
+        assert session_id.key == "test-key-123"
+
+    def test_parse_without_agent_name_uses_parsed_name(self) -> None:
+        """Test parsing @name@key format without agent_name uses name from string."""
+        session_id = AgentSessionId.parse("@ParsedAgent@test-key-123")
+
+        assert session_id.name == "ParsedAgent"
+        assert session_id.key == "test-key-123"
+
+    def test_parse_plain_string_with_agent_name(self) -> None:
+        """Test parsing plain string with agent_name uses entire string as key."""
+        session_id = AgentSessionId.parse("simple-thread-123", agent_name="TestAgent")
+
+        assert session_id.name == "TestAgent"
+        assert session_id.key == "simple-thread-123"
+
+    def test_parse_plain_string_without_agent_name_raises(self) -> None:
+        """Test parsing plain string without agent_name raises ValueError."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("simple-thread-123")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
+
+class TestDurableAgentSession:
+    """Test suite for DurableAgentSession."""
+
+    def test_init_with_durable_session_id(self) -> None:
+        """Test DurableAgentSession initialization with durable session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        session = DurableAgentSession(durable_session_id=session_id)
+
+        assert session.durable_session_id is not None
+        assert session.durable_session_id == session_id
+
+    def test_init_without_durable_session_id(self) -> None:
+        """Test DurableAgentSession initialization without durable session ID."""
+        session = DurableAgentSession()
+
+        assert session.durable_session_id is None
+
+    def test_durable_session_id_setter(self) -> None:
+        """Test setting a durable session ID to an existing session."""
+        session = DurableAgentSession()
+        assert session.durable_session_id is None
+
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        session.durable_session_id = session_id
+
+        assert session.durable_session_id is not None
+        assert session.durable_session_id == session_id
+        assert session.durable_session_id.name == "TestAgent"
+
+    def test_from_session_id(self) -> None:
+        """Test creating DurableAgentSession from session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        session = DurableAgentSession(durable_session_id=session_id)
+
+        assert isinstance(session, DurableAgentSession)
+        assert session.durable_session_id is not None
+        assert session.durable_session_id == session_id
+        assert session.durable_session_id.name == "TestAgent"
+        assert session.durable_session_id.key == "test-key"
+
+    def test_init_with_service_session_id(self) -> None:
+        """Test creating DurableAgentSession with explicit service session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        session = DurableAgentSession(durable_session_id=session_id, service_session_id="service-123")
+
+        assert session.durable_session_id is not None
+        assert session.durable_session_id == session_id
+        assert session.service_session_id == "service-123"
+
+    def test_to_dict_with_durable_session_id(self) -> None:
+        """Test serialization includes durable session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        session = DurableAgentSession(durable_session_id=session_id)
+
+        serialized = session.to_dict()
+
+        assert isinstance(serialized, dict)
+        assert "durable_session_id" in serialized
+        assert serialized["durable_session_id"] == "@TestAgent@test-key"
+
+    def test_to_dict_without_durable_session_id(self) -> None:
+        """Test serialization without durable session ID."""
+        session = DurableAgentSession()
+
+        serialized = session.to_dict()
+
+        assert isinstance(serialized, dict)
+        assert "durable_session_id" not in serialized
+
+    def test_from_dict_with_durable_session_id(self) -> None:
+        """Test deserialization restores durable session ID."""
+        serialized: dict[str, Any] = {
+            "type": "session",
+            "session_id": "session-123",
+            "service_session_id": "service-123",
+            "state": {},
+            "durable_session_id": "@TestAgent@test-key",
+        }
+
+        session = DurableAgentSession.from_dict(serialized)
+
+        assert isinstance(session, DurableAgentSession)
+        assert session.durable_session_id is not None
+        assert session.durable_session_id.name == "TestAgent"
+        assert session.durable_session_id.key == "test-key"
+        assert session.service_session_id == "service-123"
+
+    def test_from_dict_without_durable_session_id(self) -> None:
+        """Test deserialization without durable session ID."""
+        serialized: dict[str, Any] = {
+            "type": "session",
+            "session_id": "session-456",
+            "service_session_id": "service-456",
+            "state": {},
+        }
+
+        session = DurableAgentSession.from_dict(serialized)
+
+        assert isinstance(session, DurableAgentSession)
+        assert session.durable_session_id is None
+        assert session.session_id == "session-456"
+
+    def test_round_trip_serialization(self) -> None:
+        """Test round-trip serialization preserves durable session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key-789")
+        original = DurableAgentSession(durable_session_id=session_id)
+
+        serialized = original.to_dict()
+        restored = DurableAgentSession.from_dict(serialized)
+
+        assert isinstance(restored, DurableAgentSession)
+        assert restored.durable_session_id is not None
+        assert restored.durable_session_id.name == session_id.name
+        assert restored.durable_session_id.key == session_id.key
+
+    def test_from_dict_invalid_durable_session_id_type(self) -> None:
+        """Test deserialization with invalid durable session ID type raises error."""
+        serialized = {
+            "type": "session",
+            "session_id": "session-123",
+            "state": {},
+            "durable_session_id": 12345,  # Invalid type
+        }
+
+        with pytest.raises(ValueError, match="durable_session_id must be a string"):
+            DurableAgentSession.from_dict(serialized)
+
+
+class TestAgentSessionCompatibility:
+    """Test suite for compatibility between AgentSession and DurableAgentSession."""
+
+    def test_agent_session_to_dict(self) -> None:
+        """Test that base AgentSession can be serialized."""
+        session = AgentSession()
+
+        serialized = session.to_dict()
+
+        assert isinstance(serialized, dict)
+        assert "session_id" in serialized
+
+    def test_agent_session_from_dict(self) -> None:
+        """Test that base AgentSession can be deserialized."""
+        session = AgentSession()
+        serialized = session.to_dict()
+
+        restored = AgentSession.from_dict(serialized)
+
+        assert isinstance(restored, AgentSession)
+        assert restored.session_id == session.session_id
+
+    def test_durable_session_is_agent_session(self) -> None:
+        """Test that DurableAgentSession is an AgentSession."""
+        session = DurableAgentSession()
+
+        assert isinstance(session, AgentSession)
+        assert isinstance(session, DurableAgentSession)
+
+
+class TestModelIntegration:
+    """Test suite for integration between models."""
+
+    def test_session_id_string_format(self) -> None:
+        """Test that AgentSessionId string format is consistent."""
+        session_id = AgentSessionId.with_random_key("AgentEntity")
+        session_id_str = str(session_id)
+
+        assert session_id_str.startswith("@AgentEntity@")
+
+    def test_session_with_durable_id_preserves_on_serialization(self) -> None:
+        """Test that session with durable session ID preserves it through serialization."""
+        session_id = AgentSessionId(name="TestAgent", key="preserved-key")
+        session = DurableAgentSession.from_session_id(session_id)
+
+        # Serialize and deserialize
+        serialized = session.to_dict()
+        restored = DurableAgentSession.from_dict(serialized)
+
+        # Durable session ID should be preserved
+        assert restored.durable_session_id is not None
+        assert restored.durable_session_id.name == "TestAgent"
+        assert restored.durable_session_id.key == "preserved-key"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_client.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for DurableAIAgentClient.
+
+Focuses on critical client workflows: agent retrieval, protocol compliance, and integration.
+Run with: pytest tests/test_client.py -v
+"""
+
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import SupportsAgentRun
+
+from agent_framework_durabletask import DurableAgentSession, DurableAIAgentClient
+from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
+from agent_framework_durabletask._shim import DurableAIAgent
+
+
+@pytest.fixture
+def mock_grpc_client() -> Mock:
+    """Create a mock TaskHubGrpcClient for testing."""
+    return Mock()
+
+
+@pytest.fixture
+def agent_client(mock_grpc_client: Mock) -> DurableAIAgentClient:
+    """Create a DurableAIAgentClient with mock gRPC client."""
+    return DurableAIAgentClient(mock_grpc_client)
+
+
+@pytest.fixture
+def agent_client_with_custom_polling(mock_grpc_client: Mock) -> DurableAIAgentClient:
+    """Create a DurableAIAgentClient with custom polling parameters."""
+    return DurableAIAgentClient(
+        mock_grpc_client,
+        max_poll_retries=15,
+        poll_interval_seconds=0.5,
+    )
+
+
+class TestDurableAIAgentClientGetAgent:
+    """Test core workflow: retrieving agents from the client."""
+
+    def test_get_agent_returns_durable_agent_shim(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify get_agent returns a DurableAIAgent instance."""
+        agent = agent_client.get_agent("assistant")
+
+        assert isinstance(agent, DurableAIAgent)
+        assert isinstance(agent, SupportsAgentRun)  # pyrefly: ignore[unsafe-overlap]
+
+    def test_get_agent_shim_has_correct_name(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify retrieved agent has the correct name."""
+        agent = agent_client.get_agent("my_agent")
+
+        assert agent.name == "my_agent"
+
+    def test_get_agent_multiple_times_returns_new_instances(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify multiple get_agent calls return independent instances."""
+        agent1 = agent_client.get_agent("assistant")
+        agent2 = agent_client.get_agent("assistant")
+
+        assert agent1 is not agent2  # Different object instances
+
+    def test_get_agent_different_agents(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify client can retrieve multiple different agents."""
+        agent1 = agent_client.get_agent("agent1")
+        agent2 = agent_client.get_agent("agent2")
+
+        assert agent1.name == "agent1"
+        assert agent2.name == "agent2"
+
+
+class TestDurableAIAgentClientIntegration:
+    """Test integration scenarios between client and agent shim."""
+
+    def test_client_agent_has_working_run_method(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify agent from client has callable run method (even if not yet implemented)."""
+        agent = agent_client.get_agent("assistant")
+
+        assert hasattr(agent, "run")
+        assert callable(agent.run)
+
+    def test_client_agent_can_create_sessions(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify agent from client can create DurableAgentSession instances."""
+        agent = agent_client.get_agent("assistant")
+
+        session = agent.create_session()
+
+        assert isinstance(session, DurableAgentSession)
+
+
+class TestDurableAIAgentClientPollingConfiguration:
+    """Test polling configuration parameters for DurableAIAgentClient."""
+
+    def test_client_uses_default_polling_parameters(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify client initializes with default polling parameters."""
+        assert agent_client.max_poll_retries == DEFAULT_MAX_POLL_RETRIES
+        assert agent_client.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
+
+    def test_client_accepts_custom_polling_parameters(
+        self, agent_client_with_custom_polling: DurableAIAgentClient
+    ) -> None:
+        """Verify client accepts and stores custom polling parameters."""
+        assert agent_client_with_custom_polling.max_poll_retries == 15
+        assert agent_client_with_custom_polling.poll_interval_seconds == 0.5
+
+    def test_client_validates_max_poll_retries(self, mock_grpc_client: Mock) -> None:
+        """Verify client validates and normalizes max_poll_retries."""
+        # Test with zero - should enforce minimum of 1
+        client = DurableAIAgentClient(mock_grpc_client, max_poll_retries=0)
+        assert client.max_poll_retries == 1
+
+        # Test with negative - should enforce minimum of 1
+        client = DurableAIAgentClient(mock_grpc_client, max_poll_retries=-5)
+        assert client.max_poll_retries == 1
+
+    def test_client_validates_poll_interval_seconds(self, mock_grpc_client: Mock) -> None:
+        """Verify client validates and normalizes poll_interval_seconds."""
+        # Test with zero - should use default
+        client = DurableAIAgentClient(mock_grpc_client, poll_interval_seconds=0)
+        assert client.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
+
+        # Test with negative - should use default
+        client = DurableAIAgentClient(mock_grpc_client, poll_interval_seconds=-0.5)
+        assert client.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
+
+        # Test with valid float
+        client = DurableAIAgentClient(mock_grpc_client, poll_interval_seconds=2.5)
+        assert client.poll_interval_seconds == 2.5
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_durable_agent_state.py

@@ -0,0 +1,525 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for DurableAgentState and related classes."""
+
+import json
+from datetime import datetime
+
+import pytest
+from agent_framework import Content, Message, UsageDetails
+
+from agent_framework_durabletask._durable_agent_state import (
+    DurableAgentState,
+    DurableAgentStateContent,
+    DurableAgentStateFunctionCallContent,
+    DurableAgentStateMessage,
+    DurableAgentStateRequest,
+    DurableAgentStateTextContent,
+    DurableAgentStateUnknownContent,
+    DurableAgentStateUsage,
+)
+from agent_framework_durabletask._models import RunRequest
+
+
+class TestDurableAgentStateRequestOrchestrationId:
+    """Test suite for DurableAgentStateRequest orchestration_id field."""
+
+    def test_request_with_orchestration_id(self) -> None:
+        """Test creating a request with an orchestration_id."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+            orchestration_id="orch-456",
+        )
+
+        assert request.orchestration_id == "orch-456"
+
+    def test_request_to_dict_includes_orchestration_id(self) -> None:
+        """Test that to_dict includes orchestrationId when set."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+            orchestration_id="orch-789",
+        )
+
+        data = request.to_dict()
+
+        assert "orchestrationId" in data
+        assert data["orchestrationId"] == "orch-789"
+
+    def test_request_to_dict_excludes_orchestration_id_when_none(self) -> None:
+        """Test that to_dict excludes orchestrationId when not set."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+        )
+
+        data = request.to_dict()
+
+        assert "orchestrationId" not in data
+
+    def test_request_from_dict_with_orchestration_id(self) -> None:
+        """Test from_dict correctly parses orchestrationId."""
+        data = {
+            "$type": "request",
+            "correlationId": "corr-123",
+            "createdAt": "2024-01-01T00:00:00Z",
+            "messages": [{"role": "user", "contents": [{"$type": "text", "text": "test"}]}],
+            "orchestrationId": "orch-from-dict",
+        }
+
+        request = DurableAgentStateRequest.from_dict(data)
+
+        assert request.orchestration_id == "orch-from-dict"
+
+    def test_request_from_run_request_with_orchestration_id(self) -> None:
+        """Test from_run_request correctly transfers orchestration_id."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+            orchestration_id="orch-from-run-request",
+        )
+
+        durable_request = DurableAgentStateRequest.from_run_request(run_request)
+
+        assert durable_request.orchestration_id == "orch-from-run-request"
+
+    def test_request_from_run_request_without_orchestration_id(self) -> None:
+        """Test from_run_request correctly handles missing orchestration_id."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+        )
+
+        durable_request = DurableAgentStateRequest.from_run_request(run_request)
+
+        assert durable_request.orchestration_id is None
+
+
+class TestDurableAgentStateMessageCreatedAt:
+    """Test suite for DurableAgentStateMessage created_at field handling."""
+
+    def test_message_from_run_request_without_created_at_preserves_none(self) -> None:
+        """Test from_run_request handles auto-populated created_at from RunRequest.
+
+        When a RunRequest is created with None for created_at, RunRequest defaults it to
+        current UTC time. The resulting DurableAgentStateMessage should have this timestamp.
+        """
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+            created_at=None,  # RunRequest will default this to current time
+        )
+
+        durable_message = DurableAgentStateMessage.from_run_request(run_request)
+
+        # RunRequest auto-populates created_at, so it should not be None
+        assert durable_message.created_at is not None
+
+    def test_message_from_run_request_with_created_at_parses_correctly(self) -> None:
+        """Test from_run_request correctly parses a valid created_at timestamp."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+            created_at=datetime(2024, 1, 15, 10, 30, 0),
+        )
+
+        durable_message = DurableAgentStateMessage.from_run_request(run_request)
+
+        assert durable_message.created_at is not None
+        assert durable_message.created_at.year == 2024
+        assert durable_message.created_at.month == 1
+        assert durable_message.created_at.day == 15
+
+
+class TestDurableAgentState:
+    """Test suite for DurableAgentState."""
+
+    def test_schema_version(self) -> None:
+        """Test that schema version is set correctly."""
+        state = DurableAgentState()
+        assert state.schema_version == "1.1.0"
+
+    def test_to_dict_serialization(self) -> None:
+        """Test that to_dict produces correct structure."""
+        state = DurableAgentState()
+        data = state.to_dict()
+
+        assert "schemaVersion" in data
+        assert "data" in data
+        assert data["schemaVersion"] == "1.1.0"
+        assert "conversationHistory" in data["data"]
+
+    def test_from_dict_deserialization(self) -> None:
+        """Test that from_dict restores state correctly."""
+        original_data = {
+            "schemaVersion": "1.1.0",
+            "data": {
+                "conversationHistory": [
+                    {
+                        "$type": "request",
+                        "correlationId": "test-123",
+                        "createdAt": "2024-01-01T00:00:00Z",
+                        "messages": [
+                            {
+                                "role": "user",
+                                "contents": [{"$type": "text", "text": "Hello"}],
+                            }
+                        ],
+                    }
+                ]
+            },
+        }
+
+        state = DurableAgentState.from_dict(original_data)
+
+        assert state.schema_version == "1.1.0"
+        assert len(state.data.conversation_history) == 1
+        assert isinstance(state.data.conversation_history[0], DurableAgentStateRequest)
+
+    def test_round_trip_serialization(self) -> None:
+        """Test that round-trip serialization preserves data."""
+        state = DurableAgentState()
+        state.data.conversation_history.append(
+            DurableAgentStateRequest(
+                correlation_id="test-456",
+                created_at=datetime.now(),
+                messages=[
+                    DurableAgentStateMessage(
+                        role="user",
+                        contents=[DurableAgentStateTextContent(text="Test message")],
+                    )
+                ],
+            )
+        )
+
+        data = state.to_dict()
+        restored = DurableAgentState.from_dict(data)
+
+        assert restored.schema_version == state.schema_version
+        assert len(restored.data.conversation_history) == len(state.data.conversation_history)
+        assert restored.data.conversation_history[0].correlation_id == "test-456"
+
+    def test_function_call_round_trip_preserves_string_arguments(self) -> None:
+        """Function call arguments should remain strings across durable state replay."""
+        original = Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id="call-123",
+                    name="get_weather",
+                    arguments='{"location":"Chicago"}',
+                )
+            ],
+        )
+
+        durable_message = DurableAgentStateMessage.from_chat_message(original)
+        restored = durable_message.to_chat_message()
+
+        assert restored.contents[0].type == "function_call"
+        assert restored.contents[0].arguments == '{"location": "Chicago"}'
+
+    def test_function_call_content_supports_legacy_mapping_arguments(self) -> None:
+        """Existing persisted mapping arguments should still restore successfully."""
+        content = DurableAgentStateFunctionCallContent(
+            call_id="call-123",
+            name="get_weather",
+            arguments={"location": "Chicago"},
+        )
+
+        restored = content.to_ai_content()
+
+        assert restored.type == "function_call"
+        assert restored.arguments == '{"location": "Chicago"}'
+
+
+class TestDurableAgentStateUsage:
+    """Test suite for DurableAgentStateUsage."""
+
+    def test_usage_init_with_defaults(self) -> None:
+        """Test creating usage with default values."""
+        usage = DurableAgentStateUsage()
+
+        assert usage.input_token_count is None
+        assert usage.output_token_count is None
+        assert usage.total_token_count is None
+        assert usage.extensionData is None
+
+    def test_usage_init_with_values(self) -> None:
+        """Test creating usage with specific values."""
+        usage = DurableAgentStateUsage(
+            input_token_count=100,
+            output_token_count=200,
+            total_token_count=300,
+            extensionData={"custom_field": "value"},
+        )
+
+        assert usage.input_token_count == 100
+        assert usage.output_token_count == 200
+        assert usage.total_token_count == 300
+        assert usage.extensionData == {"custom_field": "value"}
+
+    def test_usage_to_dict(self) -> None:
+        """Test that to_dict produces correct structure."""
+        usage = DurableAgentStateUsage(
+            input_token_count=50,
+            output_token_count=75,
+            total_token_count=125,
+        )
+
+        data = usage.to_dict()
+
+        assert data["inputTokenCount"] == 50
+        assert data["outputTokenCount"] == 75
+        assert data["totalTokenCount"] == 125
+
+    def test_usage_to_dict_with_extension_data(self) -> None:
+        """Test that to_dict includes extensionData when present."""
+        usage = DurableAgentStateUsage(
+            input_token_count=10,
+            output_token_count=20,
+            total_token_count=30,
+            extensionData={"provider_specific": 123},
+        )
+
+        data = usage.to_dict()
+
+        assert "extensionData" in data
+        assert data["extensionData"] == {"provider_specific": 123}
+
+    def test_usage_from_dict(self) -> None:
+        """Test that from_dict restores usage correctly."""
+        data = {
+            "inputTokenCount": 100,
+            "outputTokenCount": 200,
+            "totalTokenCount": 300,
+            "extensionData": {"extra": "data"},
+        }
+
+        usage = DurableAgentStateUsage.from_dict(data)
+
+        assert usage.input_token_count == 100
+        assert usage.output_token_count == 200
+        assert usage.total_token_count == 300
+        assert usage.extensionData == {"extra": "data"}
+
+    def test_usage_from_usage_details(self) -> None:
+        """Test creating DurableAgentStateUsage from UsageDetails."""
+        usage_details: UsageDetails = {
+            "input_token_count": 150,
+            "output_token_count": 250,
+            "total_token_count": 400,
+        }
+
+        usage = DurableAgentStateUsage.from_usage(usage_details)
+
+        assert usage is not None
+        assert usage.input_token_count == 150
+        assert usage.output_token_count == 250
+        assert usage.total_token_count == 400
+
+    def test_usage_from_usage_details_with_extension_fields(self) -> None:
+        """Test that non-standard fields are captured in extensionData."""
+        usage_details: UsageDetails = {
+            "input_token_count": 100,
+            "output_token_count": 200,
+            "total_token_count": 300,
+        }
+        # Add provider-specific fields (UsageDetails is a TypedDict but allows extra keys)
+        usage_details["prompt_tokens"] = 100  # type: ignore[typeddict-unknown-key]
+        usage_details["completion_tokens"] = 200  # type: ignore[typeddict-unknown-key]
+
+        usage = DurableAgentStateUsage.from_usage(usage_details)
+
+        assert usage is not None
+        assert usage.extensionData is not None
+        assert usage.extensionData["prompt_tokens"] == 100
+        assert usage.extensionData["completion_tokens"] == 200
+
+    def test_usage_from_usage_none(self) -> None:
+        """Test that from_usage returns None for None input."""
+        usage = DurableAgentStateUsage.from_usage(None)
+
+        assert usage is None
+
+    def test_usage_to_usage_details(self) -> None:
+        """Test converting back to UsageDetails."""
+        usage = DurableAgentStateUsage(
+            input_token_count=100,
+            output_token_count=200,
+            total_token_count=300,
+        )
+
+        details = usage.to_usage_details()
+
+        assert details.get("input_token_count") == 100
+        assert details.get("output_token_count") == 200
+        assert details.get("total_token_count") == 300
+
+    def test_usage_to_usage_details_with_extension_data(self) -> None:
+        """Test that extensionData is merged into UsageDetails."""
+        usage = DurableAgentStateUsage(
+            input_token_count=50,
+            output_token_count=75,
+            total_token_count=125,
+            extensionData={"prompt_tokens": 50, "completion_tokens": 75},
+        )
+
+        details = usage.to_usage_details()
+
+        assert details.get("input_token_count") == 50
+        assert details.get("output_token_count") == 75
+        assert details.get("total_token_count") == 125
+        # Extension data should be merged into the result
+        assert details.get("prompt_tokens") == 50
+        assert details.get("completion_tokens") == 75
+
+    def test_usage_round_trip(self) -> None:
+        """Test round-trip conversion from UsageDetails to DurableAgentStateUsage and back."""
+        original: UsageDetails = {
+            "input_token_count": 100,
+            "output_token_count": 200,
+            "total_token_count": 300,
+        }
+
+        usage = DurableAgentStateUsage.from_usage(original)
+        assert usage is not None
+        restored = usage.to_usage_details()
+
+        assert restored.get("input_token_count") == original.get("input_token_count")
+        assert restored.get("output_token_count") == original.get("output_token_count")
+        assert restored.get("total_token_count") == original.get("total_token_count")
+
+
+class TestDurableAgentStateUnknownContent:
+    """Test suite for DurableAgentStateUnknownContent serialization."""
+
+    def test_unknown_content_from_content_object_produces_serializable_dict(self) -> None:
+        """Test that from_unknown_content serializes Content objects to dicts."""
+        content = Content.from_mcp_server_tool_call(
+            call_id="call-1",
+            tool_name="search",
+            server_name="learn-mcp",
+            arguments={"query": "azure functions"},
+        )
+
+        unknown = DurableAgentStateUnknownContent.from_unknown_content(content)
+        result = unknown.to_dict()
+
+        # The content field should be a dict, not a Content object
+        assert isinstance(result["content"], dict)
+        assert result["content"]["type"] == "mcp_server_tool_call"
+
+    def test_unknown_content_to_dict_is_json_serializable(self) -> None:
+        """Test that to_dict output can be passed to json.dumps without error."""
+        content = Content.from_mcp_server_tool_result(
+            call_id="call-1",
+            output="Azure Functions documentation...",
+        )
+
+        unknown = DurableAgentStateUnknownContent.from_unknown_content(content)
+        result = unknown.to_dict()
+
+        # This must not raise TypeError
+        serialized = json.dumps(result)
+        assert serialized is not None
+
+    def test_unknown_content_round_trip_preserves_content(self) -> None:
+        """Test that Content objects survive serialization and deserialization."""
+        original = Content.from_mcp_server_tool_call(
+            call_id="call-1",
+            tool_name="fetch",
+            server_name="learn-mcp",
+            arguments={"url": "https://example.com"},
+        )
+
+        unknown = DurableAgentStateUnknownContent.from_unknown_content(original)
+        restored = unknown.to_ai_content()
+
+        assert restored.type == "mcp_server_tool_call"
+        assert restored.tool_name == "fetch"
+        assert restored.server_name == "learn-mcp"
+
+    def test_unknown_content_from_plain_dict_unchanged(self) -> None:
+        """Test that non-Content values are stored as-is."""
+        plain = {"some": "data"}
+
+        unknown = DurableAgentStateUnknownContent.from_unknown_content(plain)
+
+        assert unknown.content == {"some": "data"}
+
+    def test_unknown_content_to_ai_content_fallback_on_invalid_type_dict(self) -> None:
+        """Test that to_ai_content falls back when dict has 'type' but is not valid Content."""
+        invalid = {"type": "bogus_not_a_real_content_type", "extra": "stuff"}
+        unknown = DurableAgentStateUnknownContent(content=invalid)
+
+        result = unknown.to_ai_content()
+
+        assert result.type == "unknown"
+        assert result.additional_properties == {"content": invalid}
+
+    def test_from_ai_content_unknown_type_produces_serializable_state(self) -> None:
+        """Test that unknown content types in message conversion produce JSON-serializable state."""
+        content = Content.from_mcp_server_tool_call(
+            call_id="call-1",
+            tool_name="search",
+            server_name="learn-mcp",
+            arguments={"query": "create function app"},
+        )
+
+        durable_content = DurableAgentStateContent.from_ai_content(content)
+        data = durable_content.to_dict()
+
+        # Must be fully JSON-serializable
+        serialized = json.dumps(data)
+        assert serialized is not None
+
+    def test_state_with_mcp_content_is_json_serializable(self) -> None:
+        """Test that full DurableAgentState with MCP content can be serialized to JSON.
+
+        This reproduces the scenario from issue #4719 where agent state containing
+        MCP tool content could not be serialized by Azure Durable Functions.
+        """
+        state = DurableAgentState()
+        mcp_content = Content.from_mcp_server_tool_call(
+            call_id="call-1",
+            tool_name="search",
+            server_name="learn-mcp",
+            arguments={"query": "azure functions"},
+        )
+        message = DurableAgentStateMessage.from_chat_message(Message(role="assistant", contents=[mcp_content]))
+        state.data.conversation_history.append(
+            DurableAgentStateRequest(
+                correlation_id="test-mcp",
+                created_at=datetime.now(),
+                messages=[message],
+            )
+        )
+
+        state_dict = state.to_dict()
+
+        # This simulates what Azure Durable Functions does with entity state
+        serialized = json.dumps(state_dict)
+        assert serialized is not None
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_durable_entities.py

@@ -0,0 +1,768 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for AgentEntity.
+
+Run with: pytest tests/test_entities.py -v
+"""
+
+from collections.abc import AsyncIterator
+from datetime import datetime
+from typing import Any, TypeVar
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from agent_framework import AgentResponse, AgentResponseUpdate, Content, Message, ResponseStream
+from pydantic import BaseModel
+
+from agent_framework_durabletask import (
+    AgentEntity,
+    AgentEntityStateProviderMixin,
+    DurableAgentState,
+    DurableAgentStateData,
+    DurableAgentStateMessage,
+    DurableAgentStateRequest,
+    DurableAgentStateResponse,
+    DurableAgentStateTextContent,
+    DurableAgentStateTextReasoningContent,
+    RunRequest,
+)
+from agent_framework_durabletask._entities import DurableTaskEntityStateProvider
+
+StateT = TypeVar("StateT")
+
+
+class MockEntityContext:
+    """Minimal durabletask EntityContext shim for tests."""
+
+    def __init__(self, initial_state: Any = None) -> None:
+        self._state = initial_state
+
+    def get_state(
+        self,
+        intended_type: type[StateT] | None = None,
+        default: StateT | None = None,
+    ) -> Any:
+        del intended_type
+        if self._state is None:
+            return default
+        return self._state
+
+    def set_state(self, new_state: Any) -> None:
+        self._state = new_state
+
+
+class _InMemoryStateProvider(AgentEntityStateProviderMixin):
+    """Test-only state provider for AgentEntity."""
+
+    def __init__(self, *, thread_id: str, initial_state: dict[str, Any] | None = None) -> None:
+        self._thread_id = thread_id
+        self._state_dict: dict[str, Any] = initial_state or {}
+
+    def _get_state_dict(self) -> dict[str, Any]:
+        return self._state_dict
+
+    def _set_state_dict(self, state: dict[str, Any]) -> None:
+        self._state_dict = state
+
+    def _get_thread_id_from_entity(self) -> str:
+        return self._thread_id
+
+
+def _make_entity(agent: Any, callback: Any = None, *, thread_id: str = "test-thread") -> AgentEntity:
+    return AgentEntity(agent, callback=callback, state_provider=_InMemoryStateProvider(thread_id=thread_id))
+
+
+def _role_value(chat_message: DurableAgentStateMessage) -> str:
+    """Helper to extract the string role from a Message."""
+    role = getattr(chat_message, "role", None)
+    role_value = getattr(role, "value", role)
+    if role_value is None:
+        return ""
+    return str(role_value)
+
+
+def _agent_response(text: str | None) -> AgentResponse:
+    """Create an AgentResponse with a single assistant message."""
+    message = (
+        Message(role="assistant", contents=[text]) if text is not None else Message(role="assistant", contents=[""])
+    )
+    return AgentResponse(messages=[message], created_at="2024-01-01T00:00:00Z")
+
+
+def _create_mock_run(response: AgentResponse | None = None, side_effect: Exception | None = None):
+    """Create a mock run function that handles stream parameter correctly.
+
+    The durabletask entity code tries run(stream=True) first, then falls back to run(stream=False).
+    This helper creates a mock that raises TypeError for streaming (to trigger fallback) and
+    returns the response or raises the side_effect for non-streaming.
+    """
+
+    async def mock_run(*args, stream=False, **kwargs):
+        if stream:
+            # Simulate "streaming not supported" to trigger fallback
+            raise TypeError("streaming not supported")
+        if side_effect:
+            raise side_effect
+        return response
+
+    return mock_run
+
+
+class RecordingCallback:
+    """Callback implementation capturing streaming and final responses for assertions."""
+
+    def __init__(self):
+        self.stream_mock = AsyncMock()
+        self.response_mock = AsyncMock()
+
+    async def on_streaming_response_update(
+        self,
+        update: AgentResponseUpdate,
+        context: Any,
+    ) -> None:
+        await self.stream_mock(update, context)
+
+    async def on_agent_response(self, response: AgentResponse, context: Any) -> None:
+        await self.response_mock(response, context)
+
+
+class EntityStructuredResponse(BaseModel):
+    answer: float
+
+
+class TestAgentEntityInit:
+    """Test suite for AgentEntity initialization."""
+
+    def test_init_creates_entity(self) -> None:
+        """Test that AgentEntity initializes correctly."""
+        mock_agent = Mock()
+
+        entity = _make_entity(mock_agent)
+
+        assert entity.agent == mock_agent
+        assert len(entity.state.data.conversation_history) == 0
+        assert entity.state.data.extension_data is None
+        assert entity.state.schema_version == DurableAgentState.SCHEMA_VERSION
+
+    def test_init_stores_agent_reference(self) -> None:
+        """Test that the agent reference is stored correctly."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        entity = _make_entity(mock_agent)
+
+        assert entity.agent.name == "TestAgent"
+
+    def test_init_with_different_agent_types(self) -> None:
+        """Test initialization with different agent types."""
+        agent1 = Mock()
+        agent1.__class__.__name__ = "AzureOpenAIAgent"
+
+        agent2 = Mock()
+        agent2.__class__.__name__ = "CustomAgent"
+
+        entity1 = _make_entity(agent1)
+        entity2 = _make_entity(agent2)
+
+        assert entity1.agent.__class__.__name__ == "AzureOpenAIAgent"
+        assert entity2.agent.__class__.__name__ == "CustomAgent"
+
+
+class TestDurableTaskEntityStateProvider:
+    """Tests for DurableTaskEntityStateProvider wrapper behavior and persistence wiring."""
+
+    def _make_durabletask_entity_provider(
+        self,
+        agent: Any,
+        *,
+        initial_state: dict[str, Any] | None = None,
+    ) -> tuple[DurableTaskEntityStateProvider, MockEntityContext]:
+        """Create a DurableTaskEntityStateProvider wired to an in-memory durabletask context."""
+        entity = DurableTaskEntityStateProvider()
+        ctx = MockEntityContext(initial_state)
+        # DurableEntity provides this hook; required for get_state/set_state to work in unit tests.
+        entity._initialize_entity_context(ctx)  # type: ignore[attr-defined, arg-type] # ty: ignore[invalid-argument-type]
+        return entity, ctx
+
+    def test_reset_persists_cleared_state(self) -> None:
+        mock_agent = Mock()
+
+        existing_state = {
+            "schemaVersion": "1.0.0",
+            "data": {
+                "conversationHistory": [
+                    {
+                        "$type": "request",
+                        "correlationId": "corr-existing-1",
+                        "createdAt": "2024-01-01T00:00:00Z",
+                        "messages": [{"role": "user", "contents": [{"$type": "text", "text": "msg1"}]}],
+                    }
+                ]
+            },
+        }
+
+        entity, ctx = self._make_durabletask_entity_provider(mock_agent, initial_state=existing_state)
+
+        entity.reset()
+
+        persisted = ctx.get_state(dict, default={})
+        assert isinstance(persisted, dict)
+        assert persisted["data"]["conversationHistory"] == []
+
+
+class TestAgentEntityRunAgent:
+    """Test suite for the run_agent operation."""
+
+    async def test_run_executes_agent(self) -> None:
+        """Test that run executes the agent."""
+        mock_agent = Mock()
+        mock_response = _agent_response("Test response")
+
+        # Mock run() to return response for non-streaming, raise for streaming (to test fallback)
+        async def mock_run(*args, stream=False, **kwargs):
+            if stream:
+                raise TypeError("streaming not supported")
+            return mock_response
+
+        mock_agent.run = mock_run
+
+        entity = _make_entity(mock_agent)
+
+        result = await entity.run({
+            "message": "Test message",
+            "correlationId": "corr-entity-1",
+        })
+
+        # Verify result
+        assert isinstance(result, AgentResponse)
+        assert result.text == "Test response"
+
+    async def test_run_agent_streaming_callbacks_invoked(self) -> None:
+        """Ensure streaming updates trigger callbacks when using run(stream=True)."""
+        updates = [
+            AgentResponseUpdate(contents=[Content.from_text(text="Hello")]),
+            AgentResponseUpdate(contents=[Content.from_text(text=" world")]),
+        ]
+
+        async def update_generator() -> AsyncIterator[AgentResponseUpdate]:
+            for update in updates:
+                yield update
+
+        mock_agent = Mock()
+        mock_agent.name = "StreamingAgent"
+
+        # Mock run() to return ResponseStream when stream=True
+        def mock_run(*args, stream=False, **kwargs):
+            if stream:
+                return ResponseStream(
+                    update_generator(),
+                    finalizer=AgentResponse.from_updates,
+                )
+            raise AssertionError("run(stream=False) should not be called when streaming succeeds")
+
+        mock_agent.run = mock_run
+
+        callback = RecordingCallback()
+        entity = _make_entity(mock_agent, callback=callback, thread_id="session-1")
+
+        result = await entity.run(
+            {
+                "message": "Tell me something",
+                "correlationId": "corr-stream-1",
+            },
+        )
+
+        assert isinstance(result, AgentResponse)
+        assert "Hello" in result.text
+        assert callback.stream_mock.await_count == len(updates)
+        assert callback.response_mock.await_count == 1
+
+        # Validate callback arguments
+        stream_calls = callback.stream_mock.await_args_list
+        for expected_update, recorded_call in zip(updates, stream_calls, strict=True):
+            assert recorded_call.args[0] is expected_update
+            context = recorded_call.args[1]
+            assert context.agent_name == "StreamingAgent"
+            assert context.correlation_id == "corr-stream-1"
+            assert context.thread_id == "session-1"
+            assert context.request_message == "Tell me something"
+
+        final_call = callback.response_mock.await_args
+        assert final_call is not None
+        final_response, final_context = final_call.args
+        assert final_context.agent_name == "StreamingAgent"
+        assert final_context.correlation_id == "corr-stream-1"
+        assert final_context.thread_id == "session-1"
+        assert final_context.request_message == "Tell me something"
+        assert getattr(final_response, "text", "").strip()
+
+    async def test_run_agent_final_callback_without_streaming(self) -> None:
+        """Ensure the final callback fires even when streaming is unavailable."""
+        mock_agent = Mock()
+        mock_agent.name = "NonStreamingAgent"
+        agent_response = _agent_response("Final response")
+        mock_agent.run = _create_mock_run(response=agent_response)
+
+        callback = RecordingCallback()
+        entity = _make_entity(mock_agent, callback=callback, thread_id="session-2")
+
+        result = await entity.run(
+            {
+                "message": "Hi",
+                "correlationId": "corr-final-1",
+            },
+        )
+
+        assert isinstance(result, AgentResponse)
+        assert result.text == "Final response"
+        assert callback.stream_mock.await_count == 0
+        assert callback.response_mock.await_count == 1
+
+        final_call = callback.response_mock.await_args
+        assert final_call is not None
+        assert final_call.args[0] is agent_response
+        final_context = final_call.args[1]
+        assert final_context.agent_name == "NonStreamingAgent"
+        assert final_context.correlation_id == "corr-final-1"
+        assert final_context.thread_id == "session-2"
+        assert final_context.request_message == "Hi"
+
+    async def test_run_agent_updates_conversation_history(self) -> None:
+        """Test that run_agent updates the conversation history."""
+        mock_agent = Mock()
+        mock_response = _agent_response("Agent response")
+        mock_agent.run = _create_mock_run(response=mock_response)
+
+        entity = _make_entity(mock_agent)
+
+        await entity.run({"message": "User message", "correlationId": "corr-entity-2"})
+
+        # Should have 2 entries: user message + assistant response
+        user_history = entity.state.data.conversation_history[0].messages
+        assistant_history = entity.state.data.conversation_history[1].messages
+
+        assert len(user_history) == 1
+
+        user_msg = user_history[0]
+        assert _role_value(user_msg) == "user"
+        assert user_msg.text == "User message"
+
+        assistant_msg = assistant_history[0]
+        assert _role_value(assistant_msg) == "assistant"
+        assert assistant_msg.text == "Agent response"
+
+    async def test_run_agent_increments_message_count(self) -> None:
+        """Test that run_agent increments the message count."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        assert len(entity.state.data.conversation_history) == 0
+
+        await entity.run({"message": "Message 1", "correlationId": "corr-entity-3a"})
+        assert len(entity.state.data.conversation_history) == 2
+
+        await entity.run({"message": "Message 2", "correlationId": "corr-entity-3b"})
+        assert len(entity.state.data.conversation_history) == 4
+
+        await entity.run({"message": "Message 3", "correlationId": "corr-entity-3c"})
+        assert len(entity.state.data.conversation_history) == 6
+
+    async def test_run_requires_entity_thread_id(self) -> None:
+        """Test that AgentEntity.run rejects missing entity thread identifiers."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent, thread_id="")
+
+        with pytest.raises(ValueError, match="thread_id"):
+            await entity.run({"message": "Message", "correlationId": "corr-entity-5"})
+
+    async def test_run_agent_multiple_conversations(self) -> None:
+        """Test that run_agent maintains history across multiple messages."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        # Send multiple messages
+        await entity.run({"message": "Message 1", "correlationId": "corr-entity-8a"})
+        await entity.run({"message": "Message 2", "correlationId": "corr-entity-8b"})
+        await entity.run({"message": "Message 3", "correlationId": "corr-entity-8c"})
+
+        history = entity.state.data.conversation_history
+        assert len(history) == 6
+        assert entity.state.message_count == 6
+
+    async def test_run_filters_reasoning_content_from_replayed_history(self) -> None:
+        """Replayed durable history should not include reasoning-only content items."""
+        captured_messages: list[Message] = []
+
+        async def mock_run(*args, stream=False, **kwargs):
+            if stream:
+                raise TypeError("streaming not supported")
+            captured_messages.extend(kwargs["messages"])
+            return _agent_response("Response")
+
+        mock_agent = Mock()
+        mock_agent.run = mock_run
+
+        entity = _make_entity(mock_agent)
+        entity.state.data = DurableAgentStateData(
+            conversation_history=[
+                DurableAgentStateRequest(
+                    correlation_id="corr-entity-prev-request",
+                    created_at=datetime.now(),
+                    messages=[
+                        DurableAgentStateMessage(
+                            role="user",
+                            contents=[DurableAgentStateTextContent(text="Hi")],
+                        )
+                    ],
+                ),
+                DurableAgentStateResponse(
+                    correlation_id="corr-entity-prev-response",
+                    created_at=datetime.now(),
+                    messages=[
+                        DurableAgentStateMessage(
+                            role="assistant",
+                            contents=[
+                                DurableAgentStateTextReasoningContent(text="Let me think."),
+                                DurableAgentStateTextContent(text="Hello there."),
+                            ],
+                        )
+                    ],
+                ),
+            ]
+        )
+
+        await entity.run({"message": "What next?", "correlationId": "corr-entity-replay"})
+
+        assert captured_messages
+        assert all(content.type != "reasoning" for message in captured_messages for content in message.contents)
+        assert [message.text for message in captured_messages] == ["Hi", "Hello there.", "What next?"]
+
+
+class TestAgentEntityReset:
+    """Test suite for the reset operation."""
+
+    def test_reset_clears_conversation_history(self) -> None:
+        """Test that reset clears the conversation history."""
+        mock_agent = Mock()
+        entity = _make_entity(mock_agent)
+
+        # Add some history with proper DurableAgentStateEntry objects
+        entity.state.data.conversation_history = [
+            DurableAgentStateRequest(
+                correlation_id="test-1",
+                created_at=datetime.now(),
+                messages=[
+                    DurableAgentStateMessage(
+                        role="user",
+                        contents=[DurableAgentStateTextContent(text="msg1")],
+                    )
+                ],
+            ),
+        ]
+
+        entity.reset()
+
+        assert entity.state.data.conversation_history == []
+
+    def test_reset_with_extension_data(self) -> None:
+        """Test that reset works when entity has extension data."""
+        mock_agent = Mock()
+        entity = _make_entity(mock_agent)
+
+        # Set up some initial state with conversation history
+        entity.state.data = DurableAgentStateData(conversation_history=[], extension_data={"some_key": "some_value"})
+
+        entity.reset()
+
+        assert len(entity.state.data.conversation_history) == 0
+
+    def test_reset_clears_message_count(self) -> None:
+        """Test that reset clears the message count."""
+        mock_agent = Mock()
+        entity = _make_entity(mock_agent)
+
+        entity.reset()
+
+        assert len(entity.state.data.conversation_history) == 0
+
+    async def test_reset_after_conversation(self) -> None:
+        """Test reset after a full conversation."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        # Have a conversation
+        await entity.run({"message": "Message 1", "correlationId": "corr-entity-10a"})
+        await entity.run({"message": "Message 2", "correlationId": "corr-entity-10b"})
+
+        # Verify state before reset
+        assert entity.state.message_count == 4
+        assert len(entity.state.data.conversation_history) == 4
+
+        # Reset
+        entity.reset()
+
+        # Verify state after reset
+        assert entity.state.message_count == 0
+        assert len(entity.state.data.conversation_history) == 0
+
+
+class TestErrorHandling:
+    """Test suite for error handling in entities."""
+
+    async def test_run_agent_handles_agent_exception(self) -> None:
+        """Test that run_agent handles agent exceptions."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(side_effect=Exception("Agent failed"))
+
+        entity = _make_entity(mock_agent)
+
+        result = await entity.run({"message": "Message", "correlationId": "corr-entity-error-1"})
+
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, Content)
+        assert "Agent failed" in (content.message or "")
+        assert content.error_code == "Exception"
+
+    async def test_run_agent_handles_value_error(self) -> None:
+        """Test that run_agent handles ValueError instances."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(side_effect=ValueError("Invalid input"))
+
+        entity = _make_entity(mock_agent)
+
+        result = await entity.run({"message": "Message", "correlationId": "corr-entity-error-2"})
+
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, Content)
+        assert content.error_code == "ValueError"
+        assert "Invalid input" in str(content.message)
+
+    async def test_run_agent_handles_timeout_error(self) -> None:
+        """Test that run_agent handles TimeoutError instances."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(side_effect=TimeoutError("Request timeout"))
+
+        entity = _make_entity(mock_agent)
+
+        result = await entity.run({"message": "Message", "correlationId": "corr-entity-error-3"})
+
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, Content)
+        assert content.error_code == "TimeoutError"
+
+    async def test_run_agent_preserves_message_on_error(self) -> None:
+        """Test that run_agent preserves message information on error."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(side_effect=Exception("Error"))
+
+        entity = _make_entity(mock_agent)
+
+        result = await entity.run(
+            {"message": "Test message", "correlationId": "corr-entity-error-4"},
+        )
+
+        # Even on error, message info should be preserved
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 1
+        content = result.messages[0].contents[0]
+        assert isinstance(content, Content)
+
+
+class TestConversationHistory:
+    """Test suite for conversation history tracking."""
+
+    async def test_conversation_history_has_timestamps(self) -> None:
+        """Test that conversation history entries include timestamps."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        await entity.run({"message": "Message", "correlationId": "corr-entity-history-1"})
+
+        # Check both user and assistant messages have timestamps
+        for entry in entity.state.data.conversation_history:
+            timestamp = entry.created_at
+            assert timestamp is not None
+            # Verify timestamp is in ISO format
+            datetime.fromisoformat(str(timestamp))
+
+    async def test_conversation_history_ordering(self) -> None:
+        """Test that conversation history maintains the correct order."""
+        mock_agent = Mock()
+
+        entity = _make_entity(mock_agent)
+
+        # Send multiple messages with different responses
+        mock_agent.run = _create_mock_run(response=_agent_response("Response 1"))
+        await entity.run(
+            {"message": "Message 1", "correlationId": "corr-entity-history-2a"},
+        )
+
+        mock_agent.run = _create_mock_run(response=_agent_response("Response 2"))
+        await entity.run(
+            {"message": "Message 2", "correlationId": "corr-entity-history-2b"},
+        )
+
+        mock_agent.run = _create_mock_run(response=_agent_response("Response 3"))
+        await entity.run(
+            {"message": "Message 3", "correlationId": "corr-entity-history-2c"},
+        )
+
+        # Verify order
+        history = entity.state.data.conversation_history
+        # Each conversation turn creates 2 entries: request and response
+        assert history[0].messages[0].text == "Message 1"  # Request 1
+        assert history[1].messages[0].text == "Response 1"  # Response 1
+        assert history[2].messages[0].text == "Message 2"  # Request 2
+        assert history[3].messages[0].text == "Response 2"  # Response 2
+        assert history[4].messages[0].text == "Message 3"  # Request 3
+        assert history[5].messages[0].text == "Response 3"  # Response 3
+
+    async def test_conversation_history_role_alternation(self) -> None:
+        """Test that conversation history alternates between user and assistant roles."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        await entity.run(
+            {"message": "Message 1", "correlationId": "corr-entity-history-3a"},
+        )
+        await entity.run(
+            {"message": "Message 2", "correlationId": "corr-entity-history-3b"},
+        )
+
+        # Check role alternation
+        history = entity.state.data.conversation_history
+        # Each conversation turn creates 2 entries: request and response
+        assert history[0].messages[0].role == "user"  # Request 1
+        assert history[1].messages[0].role == "assistant"  # Response 1
+        assert history[2].messages[0].role == "user"  # Request 2
+        assert history[3].messages[0].role == "assistant"  # Response 2
+
+
+class TestRunRequestSupport:
+    """Test suite for RunRequest support in entities."""
+
+    async def test_run_agent_with_run_request_object(self) -> None:
+        """Test run_agent with a RunRequest object."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        request = RunRequest(
+            message="Test message",
+            role="user",
+            enable_tool_calls=True,
+            correlation_id="corr-runreq-1",
+        )
+
+        result = await entity.run(request)
+
+        assert isinstance(result, AgentResponse)
+        assert result.text == "Response"
+
+    async def test_run_agent_with_dict_request(self) -> None:
+        """Test run_agent with a dictionary request."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        request_dict = {
+            "message": "Test message",
+            "role": "system",
+            "enable_tool_calls": False,
+            "correlationId": "corr-runreq-2",
+        }
+
+        result = await entity.run(request_dict)
+
+        assert isinstance(result, AgentResponse)
+        assert result.text == "Response"
+
+    async def test_run_agent_with_string_raises_without_correlation(self) -> None:
+        """Test that run_agent rejects legacy string input without correlation ID."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        with pytest.raises(ValueError):
+            await entity.run("Simple message")
+
+    async def test_run_agent_stores_role_in_history(self) -> None:
+        """Test that run_agent stores the role in conversation history."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        # Send as system role
+        request = RunRequest(
+            message="System message",
+            role="system",
+            correlation_id="corr-runreq-3",
+        )
+
+        await entity.run(request)
+
+        # Check that system role was stored
+        history = entity.state.data.conversation_history
+        assert history[0].messages[0].role == "system"
+        assert history[0].messages[0].text == "System message"
+
+    async def test_run_agent_with_response_format(self) -> None:
+        """Test run_agent with a JSON response format."""
+        mock_agent = Mock()
+        # Return JSON response
+        mock_agent.run = _create_mock_run(response=_agent_response('{"answer": 42}'))
+
+        entity = _make_entity(mock_agent)
+
+        request = RunRequest(
+            message="What is the answer?",
+            response_format=EntityStructuredResponse,
+            correlation_id="corr-runreq-4",
+        )
+
+        result = await entity.run(request)
+
+        assert isinstance(result, AgentResponse)
+        assert result.text == '{"answer": 42}'
+        assert result.value is None
+
+    async def test_run_agent_disable_tool_calls(self) -> None:
+        """Test run_agent with tool calls disabled."""
+        mock_agent = Mock()
+        mock_agent.run = _create_mock_run(response=_agent_response("Response"))
+
+        entity = _make_entity(mock_agent)
+
+        request = RunRequest(message="Test", enable_tool_calls=False, correlation_id="corr-runreq-5")
+
+        result = await entity.run(request)
+
+        assert isinstance(result, AgentResponse)
+        # Agent should have been called (tool disabling is framework-dependent)
+        assert result.text == "Response"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_executors.py

@@ -0,0 +1,582 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for DurableAgentExecutor implementations.
+
+Focuses on critical behavioral flows for executor strategies.
+Run with: pytest tests/test_executors.py -v
+"""
+
+import time
+from typing import Any
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import AgentResponse
+from durabletask.entities import EntityInstanceId
+from durabletask.task import Task
+from pydantic import BaseModel
+
+from agent_framework_durabletask import DurableAgentSession
+from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
+from agent_framework_durabletask._executors import (
+    ClientAgentExecutor,
+    DurableAgentTask,
+    OrchestrationAgentExecutor,
+)
+from agent_framework_durabletask._models import AgentSessionId, RunRequest
+
+
+# Fixtures
+@pytest.fixture
+def mock_client() -> Mock:
+    """Provide a mock client for ClientAgentExecutor tests."""
+    client = Mock()
+    client.signal_entity = Mock()
+    client.get_entity = Mock(return_value=None)
+    return client
+
+
+@pytest.fixture
+def mock_entity_task() -> Mock:
+    """Provide a mock entity task."""
+    task = Mock(spec=Task)
+    task.is_complete = False
+    task.is_failed = False
+    return task
+
+
+@pytest.fixture
+def mock_orchestration_context(mock_entity_task: Mock) -> Mock:
+    """Provide a mock orchestration context with call_entity configured."""
+    context = Mock()
+    context.call_entity = Mock(return_value=mock_entity_task)
+    context.new_uuid = Mock(return_value="test-uuid-1234")
+    return context
+
+
+@pytest.fixture
+def sample_run_request() -> RunRequest:
+    """Provide a sample RunRequest for tests."""
+    return RunRequest(message="test message", correlation_id="test-123")
+
+
+@pytest.fixture
+def client_executor(mock_client: Mock) -> ClientAgentExecutor:
+    """Provide a ClientAgentExecutor with minimal polling for fast tests."""
+    return ClientAgentExecutor(mock_client, max_poll_retries=1, poll_interval_seconds=0.01)
+
+
+@pytest.fixture
+def orchestration_executor(mock_orchestration_context: Mock) -> OrchestrationAgentExecutor:
+    """Provide an OrchestrationAgentExecutor."""
+    return OrchestrationAgentExecutor(mock_orchestration_context)
+
+
+@pytest.fixture
+def successful_agent_response() -> dict[str, Any]:
+    """Provide a successful agent response dictionary."""
+    return {
+        "messages": [{"role": "assistant", "contents": [{"type": "text", "text": "Hello!"}]}],
+        "created_at": "2025-12-30T10:00:00Z",
+    }
+
+
+@pytest.fixture
+def configure_successful_entity_task(mock_entity_task: Mock) -> Any:
+    """Provide a helper to configure mock_entity_task with a successful response."""
+
+    def _configure(response: dict[str, Any]) -> Mock:
+        mock_entity_task.is_failed = False
+        mock_entity_task.is_complete = False
+        mock_entity_task.get_result = Mock(return_value=response)
+        return mock_entity_task
+
+    return _configure
+
+
+@pytest.fixture
+def configure_failed_entity_task(mock_entity_task: Mock) -> Any:
+    """Provide a helper to configure mock_entity_task with a failure."""
+
+    def _configure(exception: Exception) -> Mock:
+        mock_entity_task.is_failed = True
+        mock_entity_task.is_complete = True
+        mock_entity_task.get_exception = Mock(return_value=exception)
+        return mock_entity_task
+
+    return _configure
+
+
+class TestExecutorSessionCreation:
+    """Test that executors properly create DurableAgentSession with parameters."""
+
+    def test_client_executor_creates_durable_session(self, mock_client: Mock) -> None:
+        """Verify ClientAgentExecutor creates DurableAgentSession instances."""
+        executor = ClientAgentExecutor(mock_client)
+
+        session = executor.get_new_session("test_agent")
+
+        assert isinstance(session, DurableAgentSession)
+
+    def test_client_executor_forwards_kwargs_to_session(self, mock_client: Mock) -> None:
+        """Verify ClientAgentExecutor forwards kwargs to DurableAgentSession creation."""
+        executor = ClientAgentExecutor(mock_client)
+
+        session = executor.get_new_session("test_agent", service_session_id="client-123")
+
+        assert isinstance(session, DurableAgentSession)
+        assert session.service_session_id == "client-123"
+
+    def test_orchestration_executor_creates_durable_session(
+        self, orchestration_executor: OrchestrationAgentExecutor
+    ) -> None:
+        """Verify OrchestrationAgentExecutor creates DurableAgentSession instances."""
+        session = orchestration_executor.get_new_session("test_agent")
+
+        assert isinstance(session, DurableAgentSession)
+
+    def test_orchestration_executor_forwards_kwargs_to_session(
+        self, orchestration_executor: OrchestrationAgentExecutor
+    ) -> None:
+        """Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentSession creation."""
+        session = orchestration_executor.get_new_session("test_agent", service_session_id="orch-456")
+
+        assert isinstance(session, DurableAgentSession)
+        assert session.service_session_id == "orch-456"
+
+
+class TestClientAgentExecutorRun:
+    """Test that ClientAgentExecutor.run_durable_agent works as implemented."""
+
+    def test_client_executor_run_returns_response(
+        self, client_executor: ClientAgentExecutor, sample_run_request: RunRequest
+    ) -> None:
+        """Verify ClientAgentExecutor.run_durable_agent returns AgentResponse (synchronous)."""
+        result = client_executor.run_durable_agent("test_agent", sample_run_request)
+
+        # Verify it returns an AgentResponse (synchronous, not a coroutine)
+        assert isinstance(result, AgentResponse)
+        assert result is not None
+
+
+class TestClientAgentExecutorPollingConfiguration:
+    """Test polling configuration parameters for ClientAgentExecutor."""
+
+    def test_executor_uses_default_polling_parameters(self, mock_client: Mock) -> None:
+        """Verify executor initializes with default polling parameters."""
+        executor = ClientAgentExecutor(mock_client)
+
+        assert executor.max_poll_retries == DEFAULT_MAX_POLL_RETRIES
+        assert executor.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
+
+    def test_executor_accepts_custom_polling_parameters(self, mock_client: Mock) -> None:
+        """Verify executor accepts and stores custom polling parameters."""
+        executor = ClientAgentExecutor(mock_client, max_poll_retries=20, poll_interval_seconds=0.5)
+
+        assert executor.max_poll_retries == 20
+        assert executor.poll_interval_seconds == 0.5
+
+    def test_executor_respects_custom_max_poll_retries(self, mock_client: Mock, sample_run_request: RunRequest) -> None:
+        """Verify executor respects custom max_poll_retries during polling."""
+        # Create executor with only 2 retries
+        executor = ClientAgentExecutor(mock_client, max_poll_retries=2, poll_interval_seconds=0.01)
+
+        # Run the agent
+        result = executor.run_durable_agent("test_agent", sample_run_request)
+
+        # Verify it returns AgentResponse (should timeout after 2 attempts)
+        assert isinstance(result, AgentResponse)
+
+        # Verify get_entity was called 2 times (max_poll_retries)
+        assert mock_client.get_entity.call_count == 2
+
+    def test_executor_respects_custom_poll_interval(
+        self,
+        mock_client: Mock,
+        sample_run_request: RunRequest,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Verify executor respects custom poll_interval_seconds during polling."""
+        # Create executor with very short interval
+        executor = ClientAgentExecutor(mock_client, max_poll_retries=3, poll_interval_seconds=0.01)
+
+        sleep_calls: list[float] = []
+
+        def fake_sleep(seconds: float) -> None:
+            sleep_calls.append(seconds)
+
+        # Use deterministic assertions instead of wall-clock timing to avoid CI flakiness.
+        monkeypatch.setattr("agent_framework_durabletask._executors.time.sleep", fake_sleep)
+
+        result = executor.run_durable_agent("test_agent", sample_run_request)
+
+        assert len(sleep_calls) == 3
+        assert sleep_calls == pytest.approx([0.01, 0.01, 0.01])
+        assert mock_client.get_entity.call_count == 3
+        assert isinstance(result, AgentResponse)
+
+
+class TestClientAgentExecutorFireAndForget:
+    """Test fire-and-forget mode (wait_for_response=False) for ClientAgentExecutor."""
+
+    def test_fire_and_forget_returns_immediately(self, mock_client: Mock) -> None:
+        """Verify wait_for_response=False returns immediately without polling."""
+        executor = ClientAgentExecutor(mock_client, max_poll_retries=10, poll_interval_seconds=0.1)
+
+        # Create a request with wait_for_response=False
+        request = RunRequest(message="test message", correlation_id="test-123", wait_for_response=False)
+
+        # Measure time taken
+        start = time.time()
+        result = executor.run_durable_agent("test_agent", request)
+        elapsed = time.time() - start
+
+        # Should return immediately without polling (elapsed time should be very small)
+        assert elapsed < 0.1  # Much faster than any polling would take
+
+        # Should return an AgentResponse
+        assert isinstance(result, AgentResponse)
+
+        # Should have signaled the entity but not polled
+        assert mock_client.signal_entity.call_count == 1
+        assert mock_client.get_entity.call_count == 0  # No polling occurred
+
+    def test_fire_and_forget_returns_empty_response(self, mock_client: Mock) -> None:
+        """Verify wait_for_response=False returns an acceptance message with correlation ID."""
+        executor = ClientAgentExecutor(mock_client)
+
+        request = RunRequest(message="test message", correlation_id="test-456", wait_for_response=False)
+
+        result = executor.run_durable_agent("test_agent", request)
+
+        # Verify it contains an acceptance message
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 1
+        assert result.messages[0].role == "system"
+        # Check message contains key information
+        message_text = result.messages[0].text
+        assert "accepted" in message_text.lower()
+        assert "test-456" in message_text  # Contains correlation ID
+        assert "background" in message_text.lower()
+
+
+class TestOrchestrationAgentExecutorFireAndForget:
+    """Test fire-and-forget mode for OrchestrationAgentExecutor."""
+
+    def test_orchestration_fire_and_forget_calls_signal_entity(self, mock_orchestration_context: Mock) -> None:
+        """Verify wait_for_response=False calls signal_entity instead of call_entity."""
+        executor = OrchestrationAgentExecutor(mock_orchestration_context)
+        mock_orchestration_context.signal_entity = Mock()
+
+        request = RunRequest(message="test", correlation_id="test-123", wait_for_response=False)
+
+        result = executor.run_durable_agent("test_agent", request)
+
+        # Verify signal_entity was called and call_entity was not
+        assert mock_orchestration_context.signal_entity.call_count == 1
+        assert mock_orchestration_context.call_entity.call_count == 0
+
+        # Should still return a DurableAgentTask
+        assert isinstance(result, DurableAgentTask)
+
+    def test_orchestration_fire_and_forget_returns_completed_task(self, mock_orchestration_context: Mock) -> None:
+        """Verify wait_for_response=False returns pre-completed DurableAgentTask."""
+        executor = OrchestrationAgentExecutor(mock_orchestration_context)
+        mock_orchestration_context.signal_entity = Mock()
+
+        request = RunRequest(message="test", correlation_id="test-456", wait_for_response=False)
+
+        result = executor.run_durable_agent("test_agent", request)
+
+        # Task should be immediately complete
+        assert isinstance(result, DurableAgentTask)
+        assert result.is_complete
+
+    def test_orchestration_fire_and_forget_returns_acceptance_response(self, mock_orchestration_context: Mock) -> None:
+        """Verify wait_for_response=False returns acceptance response."""
+        executor = OrchestrationAgentExecutor(mock_orchestration_context)
+        mock_orchestration_context.signal_entity = Mock()
+
+        request = RunRequest(message="test", correlation_id="test-789", wait_for_response=False)
+
+        result = executor.run_durable_agent("test_agent", request)
+
+        # Get the result
+        response = result.get_result()
+        assert isinstance(response, AgentResponse)
+        assert len(response.messages) == 1
+        assert response.messages[0].role == "system"
+        assert "test-789" in response.messages[0].text
+
+    def test_orchestration_blocking_mode_calls_call_entity(self, mock_orchestration_context: Mock) -> None:
+        """Verify wait_for_response=True uses call_entity as before."""
+        executor = OrchestrationAgentExecutor(mock_orchestration_context)
+        mock_orchestration_context.signal_entity = Mock()
+
+        request = RunRequest(message="test", correlation_id="test-abc", wait_for_response=True)
+
+        result = executor.run_durable_agent("test_agent", request)
+
+        # Verify call_entity was called and signal_entity was not
+        assert mock_orchestration_context.call_entity.call_count == 1
+        assert mock_orchestration_context.signal_entity.call_count == 0
+
+        # Should return a DurableAgentTask
+        assert isinstance(result, DurableAgentTask)
+
+
+class TestOrchestrationAgentExecutorRun:
+    """Test OrchestrationAgentExecutor.run_durable_agent implementation."""
+
+    def test_orchestration_executor_run_returns_durable_agent_task(
+        self, orchestration_executor: OrchestrationAgentExecutor, sample_run_request: RunRequest
+    ) -> None:
+        """Verify OrchestrationAgentExecutor.run_durable_agent returns DurableAgentTask."""
+        result = orchestration_executor.run_durable_agent("test_agent", sample_run_request)
+
+        assert isinstance(result, DurableAgentTask)
+
+    def test_orchestration_executor_calls_entity_with_correct_parameters(
+        self,
+        mock_orchestration_context: Mock,
+        orchestration_executor: OrchestrationAgentExecutor,
+        sample_run_request: RunRequest,
+    ) -> None:
+        """Verify call_entity is invoked with correct entity ID and request."""
+        orchestration_executor.run_durable_agent("test_agent", sample_run_request)
+
+        # Verify call_entity was called once
+        assert mock_orchestration_context.call_entity.call_count == 1
+
+        # Get the call arguments
+        call_args = mock_orchestration_context.call_entity.call_args
+        entity_id_arg = call_args[0][0]
+        operation_arg = call_args[0][1]
+        request_dict_arg = call_args[0][2]
+
+        # Verify entity ID
+        assert isinstance(entity_id_arg, EntityInstanceId)
+        assert entity_id_arg.entity == "dafx-test_agent"
+
+        # Verify operation name
+        assert operation_arg == "run"
+
+        # Verify request dict
+        assert request_dict_arg == sample_run_request.to_dict()
+
+    def test_orchestration_executor_uses_session_durable_id(
+        self,
+        mock_orchestration_context: Mock,
+        orchestration_executor: OrchestrationAgentExecutor,
+        sample_run_request: RunRequest,
+    ) -> None:
+        """Verify executor uses session's durable session ID when provided."""
+        # Create session with specific durable session ID
+        session_id = AgentSessionId(name="test_agent", key="specific-key-123")
+        session = DurableAgentSession.from_session_id(session_id)
+
+        result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, session=session)
+
+        # Verify call_entity was called with the specific key
+        call_args = mock_orchestration_context.call_entity.call_args
+        entity_id_arg = call_args[0][0]
+
+        assert entity_id_arg.key == "specific-key-123"
+        assert isinstance(result, DurableAgentTask)
+
+
+class TestDurableAgentTask:
+    """Test DurableAgentTask completion and response transformation."""
+
+    def test_durable_agent_task_transforms_successful_result(
+        self, configure_successful_entity_task: Any, successful_agent_response: dict[str, Any]
+    ) -> None:
+        """Verify DurableAgentTask converts successful entity result to AgentResponse."""
+        mock_entity_task = configure_successful_entity_task(successful_agent_response)
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
+
+        # Simulate child task completion
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        result = task.get_result()
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 1
+        assert result.messages[0].role == "assistant"
+
+    def test_durable_agent_task_propagates_failure(self, configure_failed_entity_task: Any) -> None:
+        """Verify DurableAgentTask propagates task failures."""
+        mock_entity_task = configure_failed_entity_task(ValueError("Entity error"))
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
+
+        # Simulate child task completion with failure
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        assert task.is_failed
+        # The exception is wrapped in TaskFailedError by the durabletask library
+        exception = task.get_exception()
+        assert exception is not None
+
+    def test_durable_agent_task_validates_response_format(self, configure_successful_entity_task: Any) -> None:
+        """Verify DurableAgentTask validates response format when provided."""
+        response = {
+            "messages": [{"role": "assistant", "contents": [{"type": "text", "text": '{"answer": "42"}'}]}],
+            "created_at": "2025-12-30T10:00:00Z",
+        }
+        mock_entity_task = configure_successful_entity_task(response)
+
+        class TestResponse(BaseModel):
+            answer: str
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=TestResponse, correlation_id="test-123")
+
+        # Simulate child task completion
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        result = task.get_result()
+        assert isinstance(result, AgentResponse)
+
+    def test_durable_agent_task_ignores_duplicate_completion(
+        self, configure_successful_entity_task: Any, successful_agent_response: dict[str, Any]
+    ) -> None:
+        """Verify DurableAgentTask ignores duplicate completion calls."""
+        mock_entity_task = configure_successful_entity_task(successful_agent_response)
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
+
+        # Simulate child task completion twice
+        task.on_child_completed(mock_entity_task)
+        first_result = task.get_result()
+
+        task.on_child_completed(mock_entity_task)
+        second_result = task.get_result()
+
+        # Should be the same result, get_result should only be called once
+        assert first_result is second_result
+        assert mock_entity_task.get_result.call_count == 1
+
+    def test_durable_agent_task_fails_on_malformed_response(self, configure_successful_entity_task: Any) -> None:
+        """Verify DurableAgentTask fails when entity returns malformed response data."""
+        # Use data that will cause AgentResponse.from_dict to fail
+        # Using a list instead of dict, or other invalid structure
+        mock_entity_task = configure_successful_entity_task("invalid string response")
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
+
+        # Simulate child task completion with malformed data
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        assert task.is_failed
+
+    def test_durable_agent_task_fails_on_invalid_response_format(self, configure_successful_entity_task: Any) -> None:
+        """Verify DurableAgentTask fails when response doesn't match required format."""
+        response = {
+            "messages": [{"role": "assistant", "contents": [{"type": "text", "text": '{"wrong": "field"}'}]}],
+            "created_at": "2025-12-30T10:00:00Z",
+        }
+        mock_entity_task = configure_successful_entity_task(response)
+
+        class StrictResponse(BaseModel):
+            required_field: str
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=StrictResponse, correlation_id="test-123")
+
+        # Simulate child task completion with wrong format
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        assert task.is_failed
+
+    def test_durable_agent_task_handles_empty_response(self, configure_successful_entity_task: Any) -> None:
+        """Verify DurableAgentTask handles response with empty messages list."""
+        response: dict[str, str | list[Any]] = {
+            "messages": [],
+            "created_at": "2025-12-30T10:00:00Z",
+        }
+        mock_entity_task = configure_successful_entity_task(response)
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
+
+        # Simulate child task completion
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        result = task.get_result()
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 0
+
+    def test_durable_agent_task_handles_multiple_messages(self, configure_successful_entity_task: Any) -> None:
+        """Verify DurableAgentTask correctly processes response with multiple messages."""
+        response = {
+            "messages": [
+                {"role": "assistant", "contents": [{"type": "text", "text": "First message"}]},
+                {"role": "assistant", "contents": [{"type": "text", "text": "Second message"}]},
+            ],
+            "created_at": "2025-12-30T10:00:00Z",
+        }
+        mock_entity_task = configure_successful_entity_task(response)
+
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
+
+        # Simulate child task completion
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        result = task.get_result()
+        assert isinstance(result, AgentResponse)
+        assert len(result.messages) == 2
+        assert result.messages[0].role == "assistant"
+        assert result.messages[1].role == "assistant"
+
+    def test_durable_agent_task_is_not_complete_initially(self, mock_entity_task: Mock) -> None:
+        """Verify DurableAgentTask is not complete when first created."""
+        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
+
+        assert not task.is_complete
+        assert not task.is_failed
+
+    def test_durable_agent_task_completes_with_complex_response_format(
+        self, configure_successful_entity_task: Any
+    ) -> None:
+        """Verify DurableAgentTask validates complex nested response formats correctly."""
+        response = {
+            "messages": [
+                {
+                    "role": "assistant",
+                    "contents": [
+                        {
+                            "type": "text",
+                            "text": '{"name": "test", "count": 42, "items": ["a", "b", "c"]}',
+                        }
+                    ],
+                }
+            ],
+            "created_at": "2025-12-30T10:00:00Z",
+        }
+        mock_entity_task = configure_successful_entity_task(response)
+
+        class ComplexResponse(BaseModel):
+            name: str
+            count: int
+            items: list[str]
+
+        task = DurableAgentTask(
+            entity_task=mock_entity_task, response_format=ComplexResponse, correlation_id="test-123"
+        )
+
+        # Simulate child task completion
+        task.on_child_completed(mock_entity_task)
+
+        assert task.is_complete
+        assert not task.is_failed
+        result = task.get_result()
+        assert isinstance(result, AgentResponse)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_models.py

@@ -0,0 +1,309 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for data models (RunRequest)."""
+
+import pytest
+from pydantic import BaseModel
+
+from agent_framework_durabletask._models import RunRequest
+
+
+class ModuleStructuredResponse(BaseModel):
+    value: int
+
+
+class TestRunRequest:
+    """Test suite for RunRequest."""
+
+    def test_init_with_defaults(self) -> None:
+        """Test RunRequest initialization with defaults."""
+        request = RunRequest(message="Hello", correlation_id="corr-001")
+
+        assert request.message == "Hello"
+        assert request.correlation_id == "corr-001"
+        assert request.role == "user"
+        assert request.response_format is None
+        assert request.enable_tool_calls is True
+        assert request.wait_for_response is True
+
+    def test_init_with_all_fields(self) -> None:
+        """Test RunRequest initialization with all fields."""
+        schema = ModuleStructuredResponse
+        request = RunRequest(
+            message="Hello",
+            correlation_id="corr-002",
+            role="system",
+            response_format=schema,
+            enable_tool_calls=False,
+            wait_for_response=False,
+        )
+
+        assert request.message == "Hello"
+        assert request.correlation_id == "corr-002"
+        assert request.role == "system"
+        assert request.response_format is schema
+        assert request.enable_tool_calls is False
+        assert request.wait_for_response is False
+
+    def test_init_coerces_string_role(self) -> None:
+        """Ensure string role values are coerced into Role instances."""
+        request = RunRequest(message="Hello", correlation_id="corr-003", role="system")  # type: ignore[arg-type]
+
+        assert request.role == "system"
+
+    def test_to_dict_with_defaults(self) -> None:
+        """Test to_dict with default values."""
+        request = RunRequest(message="Test message", correlation_id="corr-004")
+        data = request.to_dict()
+
+        assert data["message"] == "Test message"
+        assert data["enable_tool_calls"] is True
+        assert data["wait_for_response"] is True
+        assert data["role"] == "user"
+        assert data["correlationId"] == "corr-004"
+        assert "response_format" not in data or data["response_format"] is None
+        assert "thread_id" not in data
+
+    def test_to_dict_with_all_fields(self) -> None:
+        """Test to_dict with all fields."""
+        schema = ModuleStructuredResponse
+        request = RunRequest(
+            message="Hello",
+            correlation_id="corr-005",
+            role="assistant",
+            response_format=schema,
+            enable_tool_calls=False,
+            wait_for_response=False,
+        )
+        data = request.to_dict()
+
+        assert data["message"] == "Hello"
+        assert data["correlationId"] == "corr-005"
+        assert data["role"] == "assistant"
+        assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
+        assert data["response_format"]["module"] == schema.__module__
+        assert data["response_format"]["qualname"] == schema.__qualname__
+        assert data["enable_tool_calls"] is False
+        assert data["wait_for_response"] is False
+        assert "thread_id" not in data
+
+    def test_from_dict_with_defaults(self) -> None:
+        """Test from_dict with minimal data."""
+        data = {"message": "Hello", "correlationId": "corr-006"}
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Hello"
+        assert request.correlation_id == "corr-006"
+        assert request.role == "user"
+        assert request.enable_tool_calls is True
+        assert request.wait_for_response is True
+
+    def test_from_dict_ignores_thread_id_field(self) -> None:
+        """Ensure legacy thread_id input does not break RunRequest parsing."""
+        request = RunRequest.from_dict({"message": "Hello", "correlationId": "corr-007", "thread_id": "ignored"})
+
+        assert request.message == "Hello"
+
+    def test_from_dict_with_all_fields(self) -> None:
+        """Test from_dict with all fields."""
+        data = {
+            "message": "Test",
+            "correlationId": "corr-008",
+            "role": "system",
+            "response_format": {
+                "__response_schema_type__": "pydantic_model",
+                "module": ModuleStructuredResponse.__module__,
+                "qualname": ModuleStructuredResponse.__qualname__,
+            },
+            "enable_tool_calls": False,
+        }
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Test"
+        assert request.correlation_id == "corr-008"
+        assert request.role == "system"
+        assert request.response_format is ModuleStructuredResponse
+        assert request.enable_tool_calls is False
+
+    def test_from_dict_unknown_role_preserves_value(self) -> None:
+        """Test from_dict keeps custom roles intact."""
+        data = {"message": "Test", "correlationId": "corr-009", "role": "reviewer"}
+        request = RunRequest.from_dict(data)
+
+        assert request.role == "reviewer"
+        assert request.role != "user"
+
+    def test_from_dict_empty_message(self) -> None:
+        """Test from_dict with empty message."""
+        request = RunRequest.from_dict({"correlationId": "corr-010"})
+
+        assert request.message == ""
+        assert request.correlation_id == "corr-010"
+        assert request.role == "user"
+
+    def test_from_dict_missing_correlation_id_raises(self) -> None:
+        """Test from_dict raises when correlationId is missing."""
+        with pytest.raises(ValueError, match="correlationId is required"):
+            RunRequest.from_dict({"message": "Test"})
+
+    def test_round_trip_dict_conversion(self) -> None:
+        """Test round-trip to_dict and from_dict."""
+        original = RunRequest(
+            message="Test message",
+            correlation_id="corr-011",
+            role="system",
+            response_format=ModuleStructuredResponse,
+            enable_tool_calls=False,
+        )
+
+        data = original.to_dict()
+        restored = RunRequest.from_dict(data)
+
+        assert restored.message == original.message
+        assert restored.correlation_id == original.correlation_id
+        assert restored.role == original.role
+        assert restored.response_format is ModuleStructuredResponse
+        assert restored.enable_tool_calls == original.enable_tool_calls
+
+    def test_round_trip_with_pydantic_response_format(self) -> None:
+        """Ensure Pydantic response formats serialize and deserialize properly."""
+        original = RunRequest(
+            message="Structured",
+            correlation_id="corr-012",
+            response_format=ModuleStructuredResponse,
+        )
+
+        data = original.to_dict()
+
+        assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
+        assert data["response_format"]["module"] == ModuleStructuredResponse.__module__
+        assert data["response_format"]["qualname"] == ModuleStructuredResponse.__qualname__
+
+        restored = RunRequest.from_dict(data)
+        assert restored.response_format is ModuleStructuredResponse
+
+    def test_round_trip_with_options(self) -> None:
+        """Ensure options are preserved and response_format is deserialized."""
+        original = RunRequest(
+            message="Test",
+            correlation_id="corr-opts-1",
+            response_format=ModuleStructuredResponse,
+            enable_tool_calls=False,
+            options={
+                "response_format": ModuleStructuredResponse,
+                "enable_tool_calls": False,
+                "custom": "value",
+            },
+        )
+
+        data = original.to_dict()
+        assert data["options"]["custom"] == "value"
+
+        restored = RunRequest.from_dict(data)
+        assert restored.options is not None
+        assert restored.options["custom"] == "value"
+        assert restored.options["response_format"] is ModuleStructuredResponse
+
+    def test_init_with_correlationId(self) -> None:
+        """Test RunRequest initialization with correlationId."""
+        request = RunRequest(message="Test message", correlation_id="corr-123")
+
+        assert request.message == "Test message"
+        assert request.correlation_id == "corr-123"
+
+    def test_to_dict_with_correlationId(self) -> None:
+        """Test to_dict includes correlationId."""
+        request = RunRequest(message="Test", correlation_id="corr-456")
+        data = request.to_dict()
+
+        assert data["message"] == "Test"
+        assert data["correlationId"] == "corr-456"
+
+    def test_from_dict_with_correlationId(self) -> None:
+        """Test from_dict with correlationId."""
+        data = {"message": "Test", "correlationId": "corr-789"}
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Test"
+        assert request.correlation_id == "corr-789"
+
+    def test_round_trip_with_correlationId(self) -> None:
+        """Test round-trip to_dict and from_dict with correlationId."""
+        original = RunRequest(
+            message="Test message",
+            role="system",
+            correlation_id="corr-124",
+        )
+
+        data = original.to_dict()
+        restored = RunRequest.from_dict(data)
+
+        assert restored.message == original.message
+        assert restored.role == original.role
+        assert restored.correlation_id == original.correlation_id
+
+    def test_init_with_orchestration_id(self) -> None:
+        """Test RunRequest initialization with orchestration_id."""
+        request = RunRequest(
+            message="Test message",
+            correlation_id="corr-125",
+            orchestration_id="orch-123",
+        )
+
+        assert request.message == "Test message"
+        assert request.orchestration_id == "orch-123"
+
+    def test_to_dict_with_orchestration_id(self) -> None:
+        """Test to_dict includes orchestrationId."""
+        request = RunRequest(
+            message="Test",
+            correlation_id="corr-126",
+            orchestration_id="orch-456",
+        )
+        data = request.to_dict()
+
+        assert data["message"] == "Test"
+        assert data["orchestrationId"] == "orch-456"
+
+    def test_to_dict_excludes_orchestration_id_when_none(self) -> None:
+        """Test to_dict excludes orchestrationId when not set."""
+        request = RunRequest(
+            message="Test",
+            correlation_id="corr-127",
+        )
+        data = request.to_dict()
+
+        assert "orchestrationId" not in data
+
+    def test_from_dict_with_orchestration_id(self) -> None:
+        """Test from_dict with orchestrationId."""
+        data = {
+            "message": "Test",
+            "correlationId": "corr-128",
+            "orchestrationId": "orch-789",
+        }
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Test"
+        assert request.orchestration_id == "orch-789"
+
+    def test_round_trip_with_orchestration_id(self) -> None:
+        """Test round-trip to_dict and from_dict with orchestration_id."""
+        original = RunRequest(
+            message="Test message",
+            role="system",
+            correlation_id="corr-129",
+            orchestration_id="orch-123",
+        )
+
+        data = original.to_dict()
+        restored = RunRequest.from_dict(data)
+
+        assert restored.message == original.message
+        assert restored.role == original.role
+        assert restored.correlation_id == original.correlation_id
+        assert restored.orchestration_id == original.orchestration_id
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_orchestration_context.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for DurableAIAgentOrchestrationContext.
+
+Focuses on critical orchestration workflows: agent retrieval and integration.
+Run with: pytest tests/test_orchestration_context.py -v
+"""
+
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import SupportsAgentRun
+
+from agent_framework_durabletask import DurableAgentSession
+from agent_framework_durabletask._orchestration_context import DurableAIAgentOrchestrationContext
+from agent_framework_durabletask._shim import DurableAIAgent
+
+
+@pytest.fixture
+def mock_orchestration_context() -> Mock:
+    """Create a mock OrchestrationContext for testing."""
+    return Mock()
+
+
+@pytest.fixture
+def agent_context(mock_orchestration_context: Mock) -> DurableAIAgentOrchestrationContext:
+    """Create a DurableAIAgentOrchestrationContext with mock context."""
+    return DurableAIAgentOrchestrationContext(mock_orchestration_context)
+
+
+class TestDurableAIAgentOrchestrationContextGetAgent:
+    """Test core workflow: retrieving agents from orchestration context."""
+
+    def test_get_agent_returns_durable_agent_shim(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
+        """Verify get_agent returns a DurableAIAgent instance."""
+        agent = agent_context.get_agent("assistant")
+
+        assert isinstance(agent, DurableAIAgent)
+        assert isinstance(agent, SupportsAgentRun)  # pyrefly: ignore[unsafe-overlap]
+
+    def test_get_agent_shim_has_correct_name(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
+        """Verify retrieved agent has the correct name."""
+        agent = agent_context.get_agent("my_agent")
+
+        assert agent.name == "my_agent"
+
+    def test_get_agent_multiple_times_returns_new_instances(
+        self, agent_context: DurableAIAgentOrchestrationContext
+    ) -> None:
+        """Verify multiple get_agent calls return independent instances."""
+        agent1 = agent_context.get_agent("assistant")
+        agent2 = agent_context.get_agent("assistant")
+
+        assert agent1 is not agent2  # Different object instances
+
+    def test_get_agent_different_agents(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
+        """Verify context can retrieve multiple different agents."""
+        agent1 = agent_context.get_agent("agent1")
+        agent2 = agent_context.get_agent("agent2")
+
+        assert agent1.name == "agent1"
+        assert agent2.name == "agent2"
+
+
+class TestDurableAIAgentOrchestrationContextIntegration:
+    """Test integration scenarios between orchestration context and agent shim."""
+
+    def test_orchestration_agent_has_working_run_method(
+        self, agent_context: DurableAIAgentOrchestrationContext
+    ) -> None:
+        """Verify agent from context has callable run method (even if not yet implemented)."""
+        agent = agent_context.get_agent("assistant")
+
+        assert hasattr(agent, "run")
+        assert callable(agent.run)
+
+    def test_orchestration_agent_can_create_sessions(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
+        """Verify agent from context can create DurableAgentSession instances."""
+        agent = agent_context.get_agent("assistant")
+
+        session = agent.create_session()
+
+        assert isinstance(session, DurableAgentSession)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_shim.py

@@ -0,0 +1,228 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for DurableAIAgent shim and DurableAgentProvider.
+
+Focuses on critical message normalization, delegation, and protocol compliance.
+Run with: pytest tests/test_shim.py -v
+"""
+
+from typing import Any, cast
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import Message, SupportsAgentRun
+from pydantic import BaseModel
+
+from agent_framework_durabletask import DurableAgentSession
+from agent_framework_durabletask._executors import DurableAgentExecutor
+from agent_framework_durabletask._models import RunRequest
+from agent_framework_durabletask._shim import DurableAgentProvider, DurableAIAgent
+
+
+class ResponseFormatModel(BaseModel):
+    """Test Pydantic model for response format testing."""
+
+    result: str
+
+
+@pytest.fixture
+def mock_executor() -> Mock:
+    """Create a mock executor for testing."""
+    mock = Mock(spec=DurableAgentExecutor)
+    mock.run_durable_agent = Mock(return_value=None)
+    mock.get_new_session = Mock(return_value=DurableAgentSession())
+
+    # Mock get_run_request to create actual RunRequest objects
+    def create_run_request(
+        message: str,
+        options: dict[str, Any] | None = None,
+    ) -> RunRequest:
+        import uuid
+
+        opts = dict(options) if options else {}
+        response_format = opts.pop("response_format", None)
+        enable_tool_calls = cast("bool", opts.pop("enable_tool_calls", True))
+        wait_for_response = cast("bool", opts.pop("wait_for_response", True))
+        return RunRequest(
+            message=message,
+            correlation_id=str(uuid.uuid4()),
+            response_format=response_format,
+            enable_tool_calls=enable_tool_calls,
+            wait_for_response=wait_for_response,
+            options=opts,
+        )
+
+    mock.get_run_request = Mock(side_effect=create_run_request)
+    return mock
+
+
+@pytest.fixture
+def test_agent(mock_executor: Mock) -> DurableAIAgent[Any]:
+    """Create a test agent with mock executor."""
+    return DurableAIAgent(mock_executor, "test_agent")
+
+
+class TestDurableAIAgentMessageNormalization:
+    """Test that DurableAIAgent properly normalizes various message input types."""
+
+    def test_run_accepts_string_message(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run accepts and normalizes string messages."""
+        test_agent.run("Hello, world!")
+
+        mock_executor.run_durable_agent.assert_called_once()
+        # Verify agent_name and run_request were passed correctly as kwargs
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["agent_name"] == "test_agent"
+        assert kwargs["run_request"].message == "Hello, world!"
+
+    def test_run_accepts_chat_message(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run accepts and normalizes Message objects."""
+        chat_msg = Message(role="user", contents=["Test message"])
+        test_agent.run(chat_msg)
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["run_request"].message == "Test message"
+
+    def test_run_accepts_list_of_strings(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run accepts and joins list of strings."""
+        test_agent.run(["First message", "Second message"])
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["run_request"].message == "First message\nSecond message"
+
+    def test_run_accepts_list_of_chat_messages(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run accepts and joins list of Message objects."""
+        messages = [
+            Message(role="user", contents=["Message 1"]),
+            Message(role="assistant", contents=["Message 2"]),
+        ]
+        test_agent.run(messages)
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["run_request"].message == "Message 1\nMessage 2"
+
+    def test_run_handles_none_message(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run handles None message gracefully."""
+        test_agent.run(None)
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["run_request"].message == ""
+
+    def test_run_handles_empty_list(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run handles empty list gracefully."""
+        test_agent.run([])
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["run_request"].message == ""
+
+
+class TestDurableAIAgentParameterFlow:
+    """Test that parameters flow correctly through the shim to executor."""
+
+    def test_run_forwards_session_parameter(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run forwards session parameter to executor."""
+        session = DurableAgentSession(service_session_id="test-session")
+        test_agent.run("message", session=session)
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["session"] == session
+
+    def test_run_forwards_response_format(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run forwards response_format parameter to executor."""
+        test_agent.run("message", options={"response_format": ResponseFormatModel})
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["run_request"].response_format == ResponseFormatModel
+
+
+class TestDurableAISupportsAgentRunCompliance:
+    """Test that DurableAIAgent implements SupportsAgentRun correctly."""
+
+    def test_agent_implements_protocol(self, test_agent: DurableAIAgent[Any]) -> None:
+        """Verify DurableAIAgent implements SupportsAgentRun."""
+        assert isinstance(test_agent, SupportsAgentRun)  # pyrefly: ignore[unsafe-overlap]
+
+    def test_agent_has_required_properties(self, test_agent: DurableAIAgent[Any]) -> None:
+        """Verify DurableAIAgent has all required SupportsAgentRun properties."""
+        assert hasattr(test_agent, "id")
+        assert hasattr(test_agent, "name")
+        assert hasattr(test_agent, "display_name")
+        assert hasattr(test_agent, "description")
+
+    def test_agent_id_defaults_to_name(self, mock_executor: Mock) -> None:
+        """Verify agent id defaults to name when not provided."""
+        agent: DurableAIAgent[Any] = DurableAIAgent(mock_executor, "my_agent")
+
+        assert agent.id == "my_agent"
+        assert agent.name == "my_agent"
+
+    def test_agent_id_can_be_customized(self, mock_executor: Mock) -> None:
+        """Verify agent id can be set independently from name."""
+        agent: DurableAIAgent[Any] = DurableAIAgent(mock_executor, "my_agent", agent_id="custom-id")
+
+        assert agent.id == "custom-id"
+        assert agent.name == "my_agent"
+
+
+class TestDurableAIAgentSessionManagement:
+    """Test session creation and management."""
+
+    def test_create_session_delegates_to_executor(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify create_session delegates to executor."""
+        mock_session = DurableAgentSession()
+        mock_executor.get_new_session.return_value = mock_session
+
+        session = test_agent.create_session()
+
+        mock_executor.get_new_session.assert_called_once_with("test_agent")
+        assert session == mock_session
+
+    def test_get_session_forwards_service_session_id(
+        self, test_agent: DurableAIAgent[Any], mock_executor: Mock
+    ) -> None:
+        """Verify get_session forwards service_session_id and session_id to executor."""
+        mock_session = DurableAgentSession(service_session_id="svc-123")
+        mock_executor.get_new_session.return_value = mock_session
+
+        session = test_agent.get_session("svc-123", session_id="local-456")
+
+        mock_executor.get_new_session.assert_called_once_with(
+            "test_agent", service_session_id="svc-123", session_id="local-456"
+        )
+        assert session.service_session_id == "svc-123"
+
+    def test_get_session_without_session_id(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify get_session works with only service_session_id (session_id defaults to None)."""
+        mock_session = DurableAgentSession(service_session_id="svc-789")
+        mock_executor.get_new_session.return_value = mock_session
+
+        session = test_agent.get_session("svc-789")
+
+        mock_executor.get_new_session.assert_called_once_with(
+            "test_agent", service_session_id="svc-789", session_id=None
+        )
+        assert session.service_session_id == "svc-789"
+
+
+class TestDurableAgentProviderInterface:
+    """Test that DurableAgentProvider defines the correct interface."""
+
+    def test_provider_cannot_be_instantiated(self) -> None:
+        """Verify DurableAgentProvider is abstract and cannot be instantiated."""
+        with pytest.raises(TypeError):
+            DurableAgentProvider()  # type: ignore[abstract]
+
+    def test_provider_defines_get_agent_method(self) -> None:
+        """Verify DurableAgentProvider defines get_agent abstract method."""
+        assert hasattr(DurableAgentProvider, "get_agent")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_subworkflow_orchestration.py

@@ -0,0 +1,253 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for sub-workflow (child-orchestration) dispatch and result handling.
+
+A ``WorkflowExecutor`` node runs its inner workflow as a durable child
+orchestration. These tests cover the host-side glue:
+
+* :func:`_prepare_subworkflow_task` wraps the node's message in a trusted-input
+  marker and schedules ``dafx-{innerName}``.
+* :func:`_process_subworkflow_result` turns the child's outputs into either
+  routed messages (default) or parent outputs (``allow_direct_output``).
+* :func:`_try_unwrap_subworkflow_input` / :func:`_coerce_initial_input` reconstruct
+  the original typed object on the child side.
+"""
+
+from typing import Any
+from unittest.mock import Mock
+
+from agent_framework import WorkflowExecutor
+
+from agent_framework_durabletask._workflows.orchestrator import (
+    SUBWORKFLOW_INPUT_KEY,
+    TaskType,
+    _coerce_initial_input,
+    _prepare_subworkflow_task,
+    _process_subworkflow_result,
+    _try_unwrap_subworkflow_input,
+    _unpack_subworkflow_result,
+)
+from agent_framework_durabletask._workflows.serialization import (
+    SUBWORKFLOW_RESULT_KEY,
+    deserialize_value,
+    serialize_value,
+)
+
+
+def _subworkflow_executor(executor_id: str, inner_name: str, *, allow_direct_output: bool = False) -> Mock:
+    inner = Mock()
+    inner.name = inner_name
+    executor = Mock(spec=WorkflowExecutor)
+    executor.id = executor_id
+    executor.workflow = inner
+    executor.allow_direct_output = allow_direct_output
+    return executor
+
+
+def _event(event_type: str, executor_id: str, data: object = None) -> dict[str, Any]:
+    """Build a serialized workflow-event dict as the child orchestrator emits it."""
+    serialized: dict[str, Any] = {"type": event_type, "executor_id": executor_id}
+    if data is not None:
+        serialized["data"] = serialize_value(data)
+    return serialized
+
+
+def _result_envelope(outputs: list[Any], events: list[dict[str, Any]]) -> dict[str, Any]:
+    """Build the SUBWORKFLOW_RESULT_KEY envelope a child orchestration returns."""
+    return {SUBWORKFLOW_RESULT_KEY: True, "outputs": outputs, "events": events}
+
+
+class TestPrepareSubworkflowTask:
+    """Dispatch of a ``WorkflowExecutor`` node as a child orchestration."""
+
+    def test_schedules_inner_orchestration_by_scoped_name(self) -> None:
+        ctx = Mock()
+        ctx.call_sub_orchestrator.return_value = "task-sentinel"
+        executor = _subworkflow_executor("sub-node", "inner_wf")
+
+        task = _prepare_subworkflow_task(ctx, executor, "hello", "parent::sub-node::0")
+
+        assert task == "task-sentinel"
+        ctx.call_sub_orchestrator.assert_called_once()
+        args, kwargs = ctx.call_sub_orchestrator.call_args
+        assert args[0] == "dafx-inner_wf"
+        assert kwargs["instance_id"] == "parent::sub-node::0"
+
+    def test_wraps_message_in_marker(self) -> None:
+        ctx = Mock()
+        executor = _subworkflow_executor("sub-node", "inner_wf")
+
+        _prepare_subworkflow_task(ctx, executor, "payload", "child-id")
+
+        args, _ = ctx.call_sub_orchestrator.call_args
+        child_input = args[1]
+        # The wrapped payload round-trips back to the original message.
+        assert deserialize_value(child_input[SUBWORKFLOW_INPUT_KEY]) == "payload"
+
+
+class TestProcessSubworkflowResult:
+    """Conversion of a child orchestration's outputs into an ``ExecutorResult``."""
+
+    def test_default_routes_outputs_as_messages(self) -> None:
+        executor = _subworkflow_executor("sub-node", "inner_wf", allow_direct_output=False)
+        workflow_outputs: list[object] = []
+
+        result = _process_subworkflow_result(["a", "b"], executor, workflow_outputs)
+
+        assert result.task_type == TaskType.SUBWORKFLOW
+        assert workflow_outputs == []
+        assert result.activity_result is not None
+        sent = result.activity_result["sent_messages"]
+        assert [m["message"] for m in sent] == ["a", "b"]
+        assert all(m["source_id"] == "sub-node" and m["target_id"] is None for m in sent)
+
+    def test_allow_direct_output_extends_parent_outputs(self) -> None:
+        executor = _subworkflow_executor("sub-node", "inner_wf", allow_direct_output=True)
+        workflow_outputs: list[object] = ["existing"]
+
+        result = _process_subworkflow_result(["x", "y"], executor, workflow_outputs)
+
+        assert workflow_outputs == ["existing", "x", "y"]
+        assert result.activity_result is not None
+        assert result.activity_result["sent_messages"] == []
+
+    def test_none_result_produces_no_outputs(self) -> None:
+        executor = _subworkflow_executor("sub-node", "inner_wf")
+        workflow_outputs: list[object] = []
+
+        result = _process_subworkflow_result(None, executor, workflow_outputs)
+
+        assert result.activity_result is not None
+        assert result.activity_result["sent_messages"] == []
+        assert workflow_outputs == []
+
+    def test_scalar_result_is_wrapped_as_single_output(self) -> None:
+        executor = _subworkflow_executor("sub-node", "inner_wf", allow_direct_output=True)
+        workflow_outputs: list[object] = []
+
+        _process_subworkflow_result("solo", executor, workflow_outputs)
+
+        assert workflow_outputs == ["solo"]
+
+    def test_envelope_outputs_routed_as_messages(self) -> None:
+        """Outputs carried in a result envelope are routed like a bare-list result."""
+        executor = _subworkflow_executor("sub-node", "inner_wf", allow_direct_output=False)
+        workflow_outputs: list[object] = []
+        envelope = _result_envelope(["a", "b"], events=[])
+
+        result = _process_subworkflow_result(envelope, executor, workflow_outputs)
+
+        assert result.activity_result is not None
+        assert [m["message"] for m in result.activity_result["sent_messages"]] == ["a", "b"]
+        assert workflow_outputs == []
+
+    def test_envelope_allow_direct_output_extends_parent_outputs(self) -> None:
+        executor = _subworkflow_executor("sub-node", "inner_wf", allow_direct_output=True)
+        workflow_outputs: list[object] = []
+        envelope = _result_envelope(["x", "y"], events=[])
+
+        _process_subworkflow_result(envelope, executor, workflow_outputs)
+
+        assert workflow_outputs == ["x", "y"]
+
+    def test_intermediate_events_bubbled_retagged_with_node_id(self) -> None:
+        """A child's intermediate events bubble up re-tagged with the node id.
+
+        Mirrors the in-process WorkflowExecutor, which forwards child intermediate
+        emissions as WorkflowEvent("intermediate", executor_id=self.id, ...) so an
+        outer observer sees nested progress without the child's internal layout.
+        """
+        executor = _subworkflow_executor("sub-node", "inner_wf")
+        workflow_outputs: list[object] = []
+        envelope = _result_envelope(
+            outputs=["out"],
+            events=[_event("intermediate", "inner-exec", data="progress")],
+        )
+
+        result = _process_subworkflow_result(envelope, executor, workflow_outputs)
+
+        assert result.activity_result is not None
+        bubbled = result.activity_result["events"]
+        assert len(bubbled) == 1
+        # Re-tagged with the WorkflowExecutor node id, not the child's executor id.
+        assert bubbled[0]["executor_id"] == "sub-node"
+        assert bubbled[0]["type"] == "intermediate"
+        # Payload is preserved (still serialized for the parent timeline).
+        assert deserialize_value(bubbled[0]["data"]) == "progress"
+
+    def test_non_intermediate_child_events_are_not_bubbled(self) -> None:
+        """Only intermediate events bubble: lifecycle/output events stay child-internal."""
+        executor = _subworkflow_executor("sub-node", "inner_wf")
+        workflow_outputs: list[object] = []
+        envelope = _result_envelope(
+            outputs=["out"],
+            events=[
+                _event("executor_invoked", "inner-exec"),
+                _event("executor_completed", "inner-exec"),
+                _event("output", "inner-exec", data="out"),
+            ],
+        )
+
+        result = _process_subworkflow_result(envelope, executor, workflow_outputs)
+
+        assert result.activity_result is not None
+        assert result.activity_result["events"] == []
+
+
+class TestUnpackSubworkflowResult:
+    """Splitting a child orchestration's return value into ``(outputs, events)``."""
+
+    def test_unpacks_result_envelope(self) -> None:
+        events = [_event("intermediate", "inner-exec", data="p")]
+        envelope = _result_envelope(["a", "b"], events=events)
+
+        outputs, parsed_events = _unpack_subworkflow_result(envelope)
+
+        assert outputs == ["a", "b"]
+        assert parsed_events == events
+
+    def test_bare_list_is_outputs_with_no_events(self) -> None:
+        assert _unpack_subworkflow_result(["a", "b"]) == (["a", "b"], [])
+
+    def test_none_is_empty_outputs_and_events(self) -> None:
+        assert _unpack_subworkflow_result(None) == ([], [])
+
+    def test_scalar_is_single_output(self) -> None:
+        assert _unpack_subworkflow_result("solo") == (["solo"], [])
+
+    def test_envelope_with_missing_keys_degrades_gracefully(self) -> None:
+        """A malformed envelope (missing outputs/events) yields empty lists, not errors."""
+        outputs, events = _unpack_subworkflow_result({SUBWORKFLOW_RESULT_KEY: True})
+
+        assert outputs == []
+        assert events == []
+
+
+class TestSubworkflowInputUnwrap:
+    """Child-side reconstruction of the parent-supplied marker payload."""
+
+    def test_unwrap_detects_and_reconstructs_marker(self) -> None:
+        marker = {SUBWORKFLOW_INPUT_KEY: "wrapped"}
+
+        unwrapped, inner = _try_unwrap_subworkflow_input(marker)
+
+        assert unwrapped is True
+        assert inner == "wrapped"
+
+    def test_unwrap_ignores_non_marker_dict(self) -> None:
+        unwrapped, inner = _try_unwrap_subworkflow_input({"some": "data"})
+
+        assert unwrapped is False
+        assert inner is None
+
+    def test_unwrap_ignores_non_dict(self) -> None:
+        assert _try_unwrap_subworkflow_input("plain") == (False, None)
+
+    def test_coerce_initial_input_returns_unwrapped_inner(self) -> None:
+        # When the workflow runs as a child, _coerce_initial_input returns the
+        # reconstructed inner object directly, bypassing start-executor coercion.
+        workflow = Mock()
+        workflow.executors = {}
+        marker = {SUBWORKFLOW_INPUT_KEY: "inner-message"}
+
+        assert _coerce_initial_input(workflow, marker) == "inner-message"

python/packages/durabletask/tests/test_worker.py

@@ -0,0 +1,473 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for DurableAIAgentWorker.
+
+Focuses on critical worker flows: agent registration, validation, callbacks, and lifecycle.
+"""
+
+from unittest.mock import Mock
+
+import pytest
+
+from agent_framework_durabletask import DurableAIAgentWorker
+
+
+@pytest.fixture
+def mock_grpc_worker() -> Mock:
+    """Create a mock TaskHubGrpcWorker for testing."""
+    mock = Mock()
+    mock.add_entity = Mock(return_value="dafx-test_agent")
+    mock.start = Mock()
+    mock.stop = Mock()
+    return mock
+
+
+@pytest.fixture
+def mock_agent() -> Mock:
+    """Create a mock agent for testing."""
+    agent = Mock()
+    agent.name = "test_agent"
+    return agent
+
+
+@pytest.fixture
+def agent_worker(mock_grpc_worker: Mock) -> DurableAIAgentWorker:
+    """Create a DurableAIAgentWorker with mock worker."""
+    return DurableAIAgentWorker(mock_grpc_worker)
+
+
+class TestDurableAIAgentWorkerRegistration:
+    """Test agent registration behavior."""
+
+    def test_add_agent_accepts_agent_with_name(
+        self, agent_worker: DurableAIAgentWorker, mock_agent: Mock, mock_grpc_worker: Mock
+    ) -> None:
+        """Verify that agents with names can be registered."""
+        agent_worker.add_agent(mock_agent)
+
+        # Verify entity was registered with underlying worker
+        mock_grpc_worker.add_entity.assert_called_once()
+        # Verify agent name is tracked
+        assert "test_agent" in agent_worker.registered_agent_names
+
+    def test_add_agent_rejects_agent_without_name(self, agent_worker: DurableAIAgentWorker) -> None:
+        """Verify that agents without names are rejected."""
+        agent_no_name = Mock()
+        agent_no_name.name = None
+
+        with pytest.raises(ValueError, match="Agent must have a name"):
+            agent_worker.add_agent(agent_no_name)
+
+    def test_add_agent_rejects_empty_name(self, agent_worker: DurableAIAgentWorker) -> None:
+        """Verify that agents with empty names are rejected."""
+        agent_empty_name = Mock()
+        agent_empty_name.name = ""
+
+        with pytest.raises(ValueError, match="Agent must have a name"):
+            agent_worker.add_agent(agent_empty_name)
+
+    def test_add_agent_rejects_duplicate_names(self, agent_worker: DurableAIAgentWorker, mock_agent: Mock) -> None:
+        """Verify duplicate agent names are not allowed."""
+        agent_worker.add_agent(mock_agent)
+
+        # Try to register another agent with the same name
+        duplicate_agent = Mock()
+        duplicate_agent.name = "test_agent"
+
+        with pytest.raises(ValueError, match="already registered"):
+            agent_worker.add_agent(duplicate_agent)
+
+    def test_registered_agent_names_tracks_multiple_agents(self, agent_worker: DurableAIAgentWorker) -> None:
+        """Verify registered_agent_names tracks all registered agents."""
+        agent1 = Mock()
+        agent1.name = "agent1"
+        agent2 = Mock()
+        agent2.name = "agent2"
+        agent3 = Mock()
+        agent3.name = "agent3"
+
+        agent_worker.add_agent(agent1)
+        agent_worker.add_agent(agent2)
+        agent_worker.add_agent(agent3)
+
+        registered = agent_worker.registered_agent_names
+        assert "agent1" in registered
+        assert "agent2" in registered
+        assert "agent3" in registered
+        assert len(registered) == 3
+
+
+class TestDurableAIAgentWorkerCallbacks:
+    """Test callback configuration behavior."""
+
+    def test_worker_level_callback_accepted(self, mock_grpc_worker: Mock) -> None:
+        """Verify worker-level callback can be set."""
+        mock_callback = Mock()
+        agent_worker = DurableAIAgentWorker(mock_grpc_worker, callback=mock_callback)
+
+        assert agent_worker is not None
+
+    def test_agent_level_callback_accepted(self, agent_worker: DurableAIAgentWorker, mock_agent: Mock) -> None:
+        """Verify agent-level callback can be set during registration."""
+        mock_callback = Mock()
+
+        # Should not raise exception
+        agent_worker.add_agent(mock_agent, callback=mock_callback)
+
+        assert "test_agent" in agent_worker.registered_agent_names
+
+    def test_none_callback_accepted(self, mock_grpc_worker: Mock, mock_agent: Mock) -> None:
+        """Verify None callback is valid (no callbacks required)."""
+        agent_worker = DurableAIAgentWorker(mock_grpc_worker, callback=None)
+        agent_worker.add_agent(mock_agent, callback=None)
+
+        assert "test_agent" in agent_worker.registered_agent_names
+
+
+class TestDurableAIAgentWorkerLifecycle:
+    """Test worker lifecycle behavior."""
+
+    def test_start_delegates_to_underlying_worker(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """Verify start() delegates to wrapped worker."""
+        agent_worker.start()
+
+        mock_grpc_worker.start.assert_called_once()
+
+    def test_stop_delegates_to_underlying_worker(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """Verify stop() delegates to wrapped worker."""
+        agent_worker.stop()
+
+        mock_grpc_worker.stop.assert_called_once()
+
+    def test_start_works_with_no_agents(self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock) -> None:
+        """Verify worker can start even with no agents registered."""
+        agent_worker.start()
+
+        mock_grpc_worker.start.assert_called_once()
+
+    def test_start_works_with_multiple_agents(self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock) -> None:
+        """Verify worker can start with multiple agents registered."""
+        agent1 = Mock()
+        agent1.name = "agent1"
+        agent2 = Mock()
+        agent2.name = "agent2"
+
+        agent_worker.add_agent(agent1)
+        agent_worker.add_agent(agent2)
+        agent_worker.start()
+
+        mock_grpc_worker.start.assert_called_once()
+        assert len(agent_worker.registered_agent_names) == 2
+
+
+class TestDurableAIAgentWorkerWorkflow:
+    """Test workflow registration, including the agent-executor identity fix."""
+
+    def test_add_agent_with_entity_id_registers_under_override(
+        self, agent_worker: DurableAIAgentWorker, mock_agent: Mock
+    ) -> None:
+        """An explicit entity_id overrides the agent name as the entity identity."""
+        agent_worker.add_agent(mock_agent, entity_id="node-7")
+
+        assert "node-7" in agent_worker.registered_agent_names
+        assert "test_agent" not in agent_worker.registered_agent_names
+
+    def test_configure_workflow_registers_agent_entity_by_executor_id(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """Workflow agent executors register entities keyed by the workflow-scoped id.
+
+        The orchestrator dispatches by the scoped identity
+        ``{workflow}-{executorId}``, so an ``AgentExecutor(agent, id=...)`` whose id
+        differs from the agent name must still be reachable under that scoped id.
+        """
+        from agent_framework import AgentExecutor
+
+        agent = Mock()
+        agent.name = "Reviewer"
+        agent_executor = Mock(spec=AgentExecutor)
+        agent_executor.id = "custom-executor-id"
+        agent_executor.agent = agent
+
+        workflow = Mock()
+        workflow.name = "review"
+        workflow.executors = {"custom-executor-id": agent_executor}
+
+        agent_worker.configure_workflow(workflow)
+
+        assert "review-custom-executor-id" in agent_worker.registered_agent_names
+        assert "Reviewer" not in agent_worker.registered_agent_names
+        assert "custom-executor-id" not in agent_worker.registered_agent_names
+        mock_grpc_worker.add_orchestrator.assert_called_once()
+
+    def test_configure_workflow_registers_non_agent_executor_as_activity(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """Non-agent executors are registered as activities, not entities."""
+        from agent_framework import Executor
+
+        activity_executor = Mock(spec=Executor)
+        activity_executor.id = "router-node"
+
+        workflow = Mock()
+        workflow.name = "route"
+        workflow.executors = {"router-node": activity_executor}
+
+        agent_worker.configure_workflow(workflow)
+
+        assert agent_worker.registered_agent_names == []
+        mock_grpc_worker.add_activity.assert_called_once()
+        mock_grpc_worker.add_orchestrator.assert_called_once()
+        # The activity is registered under the workflow-scoped name.
+        registered_activity = mock_grpc_worker.add_activity.call_args[0][0]
+        assert registered_activity.__name__ == "dafx-route-router-node"
+
+
+class TestMultiWorkflowRegistration:
+    """Test hosting multiple workflows on one worker with scoped names."""
+
+    def _agent_workflow(self, name: str, executor_id: str) -> Mock:
+        from agent_framework import AgentExecutor
+
+        agent = Mock()
+        agent.name = "Assistant"
+        agent_executor = Mock(spec=AgentExecutor)
+        agent_executor.id = executor_id
+        agent_executor.agent = agent
+
+        workflow = Mock()
+        workflow.name = name
+        workflow.executors = {executor_id: agent_executor}
+        return workflow
+
+    def test_two_workflows_reusing_executor_id_do_not_collide(self, agent_worker: DurableAIAgentWorker) -> None:
+        """Two workflows that reuse an executor id register distinct scoped entities."""
+        agent_worker.configure_workflow(self._agent_workflow("orders", "assistant"))
+        agent_worker.configure_workflow(self._agent_workflow("billing", "assistant"))
+
+        assert "orders-assistant" in agent_worker.registered_agent_names
+        assert "billing-assistant" in agent_worker.registered_agent_names
+        assert set(agent_worker.registered_workflow_names) == {"orders", "billing"}
+
+    def test_registers_one_orchestrator_per_workflow(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """Each configured workflow registers its own orchestrator."""
+        agent_worker.configure_workflow(self._agent_workflow("orders", "a"))
+        agent_worker.configure_workflow(self._agent_workflow("billing", "b"))
+
+        assert mock_grpc_worker.add_orchestrator.call_count == 2
+        registered_names = {call.args[0].__name__ for call in mock_grpc_worker.add_orchestrator.call_args_list}
+        assert registered_names == {"dafx-orders", "dafx-billing"}
+
+    def test_rejects_duplicate_workflow_name(self, agent_worker: DurableAIAgentWorker) -> None:
+        """Configuring two workflows with the same name is rejected."""
+        agent_worker.configure_workflow(self._agent_workflow("orders", "a"))
+
+        with pytest.raises(ValueError, match="already registered"):
+            agent_worker.configure_workflow(self._agent_workflow("orders", "b"))
+
+    def test_rejects_case_insensitive_duplicate_workflow_name(self, agent_worker: DurableAIAgentWorker) -> None:
+        """Workflow names that differ only by case collide and are rejected.
+
+        The route ownership guard folds case, so allowing both ``orders`` and
+        ``Orders`` would let one workflow's surface reach the other's instances.
+        """
+        agent_worker.configure_workflow(self._agent_workflow("orders", "a"))
+
+        with pytest.raises(ValueError, match="case-insensitively"):
+            agent_worker.configure_workflow(self._agent_workflow("Orders", "b"))
+
+    def test_rejects_auto_generated_workflow_name(self, agent_worker: DurableAIAgentWorker) -> None:
+        """A workflow with an auto-generated WorkflowBuilder name is rejected."""
+        import uuid
+
+        workflow = self._agent_workflow(f"WorkflowBuilder-{uuid.uuid4()}", "a")
+
+        with pytest.raises(ValueError, match="auto-generated"):
+            agent_worker.configure_workflow(workflow)
+
+    def test_rejects_invalid_workflow_name(self, agent_worker: DurableAIAgentWorker) -> None:
+        """A workflow with an invalid name is rejected."""
+        workflow = self._agent_workflow("has space", "a")
+
+        with pytest.raises(ValueError, match="invalid"):
+            agent_worker.configure_workflow(workflow)
+
+
+class TestSubworkflowRegistration:
+    """Test recursive registration of nested sub-workflows on one worker."""
+
+    def _inner_agent_workflow(self, name: str, executor_id: str) -> Mock:
+        from agent_framework import AgentExecutor
+
+        agent = Mock()
+        agent.name = "InnerAssistant"
+        agent_executor = Mock(spec=AgentExecutor)
+        agent_executor.id = executor_id
+        agent_executor.agent = agent
+
+        workflow = Mock()
+        workflow.name = name
+        workflow.executors = {executor_id: agent_executor}
+        return workflow
+
+    def _outer_workflow(self, name: str, inner: Mock, *, sub_ids: tuple[str, ...] = ("sub",)) -> Mock:
+        from agent_framework import Executor, WorkflowExecutor
+
+        executors: dict[str, Mock] = {}
+        for sub_id in sub_ids:
+            sub = Mock(spec=WorkflowExecutor)
+            sub.id = sub_id
+            sub.workflow = inner
+            sub.allow_direct_output = False
+            executors[sub_id] = sub
+
+        router = Mock(spec=Executor)
+        router.id = "router"
+        executors["router"] = router
+
+        workflow = Mock()
+        workflow.name = name
+        workflow.executors = executors
+        return workflow
+
+    def test_nested_workflow_registers_both_orchestrations(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """Configuring an outer workflow registers the inner workflow's orchestration too."""
+        inner = self._inner_agent_workflow("inner", "agent_node")
+        outer = self._outer_workflow("outer", inner)
+
+        agent_worker.configure_workflow(outer)
+
+        registered = {call.args[0].__name__ for call in mock_grpc_worker.add_orchestrator.call_args_list}
+        assert registered == {"dafx-outer", "dafx-inner"}
+
+    def test_nested_workflow_registers_inner_agent_scoped(self, agent_worker: DurableAIAgentWorker) -> None:
+        """The inner workflow's agent is registered under the inner-scoped id."""
+        inner = self._inner_agent_workflow("inner", "agent_node")
+        outer = self._outer_workflow("outer", inner)
+
+        agent_worker.configure_workflow(outer)
+
+        assert "inner-agent_node" in agent_worker.registered_agent_names
+
+    def test_subworkflow_node_not_registered_as_activity(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """A WorkflowExecutor node is driven as a child orchestration, not an activity."""
+        inner = self._inner_agent_workflow("inner", "agent_node")
+        outer = self._outer_workflow("outer", inner)
+
+        agent_worker.configure_workflow(outer)
+
+        # Only the outer 'router' non-agent executor becomes an activity.
+        registered_activities = {call.args[0].__name__ for call in mock_grpc_worker.add_activity.call_args_list}
+        assert registered_activities == {"dafx-outer-router"}
+
+    def test_top_level_names_exclude_nested_workflows(self, agent_worker: DurableAIAgentWorker) -> None:
+        """``registered_workflow_names`` reports only top-level workflows."""
+        inner = self._inner_agent_workflow("inner", "agent_node")
+        outer = self._outer_workflow("outer", inner)
+
+        agent_worker.configure_workflow(outer)
+
+        assert agent_worker.registered_workflow_names == ["outer"]
+
+    def test_shared_subworkflow_registered_once(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """A sub-workflow reused by two nodes registers its orchestration only once."""
+        inner = self._inner_agent_workflow("inner", "agent_node")
+        outer = self._outer_workflow("outer", inner, sub_ids=("sub_a", "sub_b"))
+
+        agent_worker.configure_workflow(outer)
+
+        registered = [call.args[0].__name__ for call in mock_grpc_worker.add_orchestrator.call_args_list]
+        assert sorted(registered) == ["dafx-inner", "dafx-outer"]
+
+    def test_nested_workflow_with_invalid_name_is_rejected(self, agent_worker: DurableAIAgentWorker) -> None:
+        """A nested sub-workflow must also have a valid, stable name."""
+        inner = self._inner_agent_workflow("has space", "agent_node")
+        outer = self._outer_workflow("outer", inner)
+
+        with pytest.raises(ValueError, match="invalid"):
+            agent_worker.configure_workflow(outer)
+
+    def test_different_subworkflow_sharing_a_name_is_rejected(self, agent_worker: DurableAIAgentWorker) -> None:
+        """Two different sub-workflow instances that share a name collide and are rejected."""
+        from agent_framework import Executor, WorkflowExecutor
+
+        inner_a = self._inner_agent_workflow("shared", "agent_node")
+        inner_b = self._inner_agent_workflow("shared", "other_node")  # different instance, same name
+
+        sub_a = Mock(spec=WorkflowExecutor)
+        sub_a.id = "a"
+        sub_a.workflow = inner_a
+        sub_b = Mock(spec=WorkflowExecutor)
+        sub_b.id = "b"
+        sub_b.workflow = inner_b
+        router = Mock(spec=Executor)
+        router.id = "router"
+        outer = Mock()
+        outer.name = "outer"
+        outer.executors = {"a": sub_a, "b": sub_b, "router": router}
+
+        with pytest.raises(ValueError, match="different workflow|different workflows"):
+            agent_worker.configure_workflow(outer)
+
+    def test_cross_registration_nested_collision_is_atomic(
+        self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
+    ) -> None:
+        """A later configure_workflow whose nested child collides leaves the worker unchanged.
+
+        Reproduces the partial-registration path: configure one workflow, then configure
+        a second whose nested sub-workflow reuses the first's child name. The second call
+        must raise *before* mutating any state, so the second top-level workflow is not
+        left half-registered (which would also wedge a corrected retry on the duplicate
+        guard).
+        """
+        shared_a = self._inner_agent_workflow("shared", "agent_node")
+        agent_worker.configure_workflow(self._outer_workflow("first", shared_a))
+
+        orchestrators_before = mock_grpc_worker.add_orchestrator.call_count
+
+        # A *different* 'shared' instance nested under a new top-level workflow collides.
+        shared_b = self._inner_agent_workflow("shared", "other_node")
+        with pytest.raises(ValueError, match="collides"):
+            agent_worker.configure_workflow(self._outer_workflow("second", shared_b))
+
+        # The worker is not partially configured: 'second' was never added, and no new
+        # orchestration was registered.
+        assert agent_worker.registered_workflow_names == ["first"]
+        assert mock_grpc_worker.add_orchestrator.call_count == orchestrators_before
+
+    def test_executor_id_with_reserved_separator_is_rejected(self, agent_worker: DurableAIAgentWorker) -> None:
+        """An executor id containing the nested-HITL separator is rejected at registration."""
+        workflow = self._agent_workflow_with_executor_id("orders", "bad~id")
+
+        with pytest.raises(ValueError, match="reserved sub-workflow request separator"):
+            agent_worker.configure_workflow(workflow)
+
+    @staticmethod
+    def _agent_workflow_with_executor_id(name: str, executor_id: str) -> Mock:
+        from agent_framework import AgentExecutor
+
+        agent = Mock()
+        agent.name = "Assistant"
+        agent_executor = Mock(spec=AgentExecutor)
+        agent_executor.id = executor_id
+        agent_executor.agent = agent
+        workflow = Mock()
+        workflow.name = name
+        workflow.executors = {executor_id: agent_executor}
+        return workflow
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_workflow_activity.py

@@ -0,0 +1,123 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for execute_workflow_activity (shared non-agent executor activity body).
+
+These tests exercise the host-agnostic activity execution shared by the Azure
+Functions and standalone durabletask workflow hosts. In particular they protect
+the state snapshot/diff semantics: the snapshot must be a *deep* copy so that
+in-place mutations to nested objects (dicts, lists) are correctly detected as
+updates (regression guard for the shallow-copy bug, #4500).
+"""
+
+import json
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+from agent_framework_durabletask import execute_workflow_activity
+from agent_framework_durabletask._workflows.orchestrator import SOURCE_ORCHESTRATOR
+
+
+def _make_executor(executor_id: str, mutate: Any) -> Mock:
+    """Build a mock non-agent executor whose execute() mutates shared state."""
+    executor = Mock()
+    executor.id = executor_id
+    executor.execute = AsyncMock(side_effect=mutate)
+    return executor
+
+
+def _run(executor: Mock, snapshot: dict[str, Any]) -> dict[str, Any]:
+    """Invoke execute_workflow_activity and return the parsed result dict."""
+    input_data = json.dumps({
+        "message": "test",
+        "shared_state_snapshot": snapshot,
+        "source_executor_ids": [SOURCE_ORCHESTRATOR],
+    })
+    return json.loads(execute_workflow_activity(executor, input_data))
+
+
+class TestExecuteWorkflowActivityStateDiff:
+    """State snapshot/diff behavior of the shared workflow activity body."""
+
+    def test_nested_dict_mutation_detected(self) -> None:
+        """In-place mutation of a nested dict is reported as an update."""
+
+        async def mutate(message: Any, source_executor_ids: Any, state: Any, runner_context: Any) -> None:
+            config = state.get("Local.config")
+            config["code"] = "SOMECODEXXX"
+            config["enabled"] = True
+            state.commit()
+
+        executor = _make_executor("test-exec", mutate)
+        result = _run(executor, {"Local.config": {"code": "", "enabled": False}, "simple_key": "simple_value"})
+
+        updates = result["shared_state_updates"]
+        assert "Local.config" in updates, "nested mutation not detected — snapshot may be a shallow copy"
+        assert updates["Local.config"]["code"] == "SOMECODEXXX"
+        assert updates["Local.config"]["enabled"] is True
+
+    def test_new_key_in_nested_dict_detected(self) -> None:
+        """Adding a key to a nested dict is reported as an update."""
+
+        async def mutate(message: Any, source_executor_ids: Any, state: Any, runner_context: Any) -> None:
+            state.get("Local.data")["code"] = "NEW_CODE"
+            state.commit()
+
+        executor = _make_executor("test-exec", mutate)
+        result = _run(executor, {"Local.data": {"existing": "value"}})
+
+        assert result["shared_state_updates"]["Local.data"]["code"] == "NEW_CODE"
+
+    def test_nested_list_mutation_detected(self) -> None:
+        """Appending to a nested list is reported as an update."""
+
+        async def mutate(message: Any, source_executor_ids: Any, state: Any, runner_context: Any) -> None:
+            state.get("Local.items").append(4)
+            state.commit()
+
+        executor = _make_executor("test-exec", mutate)
+        result = _run(executor, {"Local.items": [1, 2, 3]})
+
+        assert result["shared_state_updates"]["Local.items"] == [1, 2, 3, 4]
+
+    def test_new_top_level_key_detected(self) -> None:
+        """Setting a new top-level key is reported as an update."""
+
+        async def mutate(message: Any, source_executor_ids: Any, state: Any, runner_context: Any) -> None:
+            state.set("Local.code", "SOMECODEXXX")
+            state.commit()
+
+        executor = _make_executor("test-exec", mutate)
+        result = _run(executor, {"existing": "value"})
+
+        assert result["shared_state_updates"]["Local.code"] == "SOMECODEXXX"
+
+    def test_unchanged_state_produces_empty_diff(self) -> None:
+        """Unmodified state produces no updates."""
+
+        async def mutate(message: Any, source_executor_ids: Any, state: Any, runner_context: Any) -> None:
+            # No mutations performed.
+            state.commit()
+
+        executor = _make_executor("test-exec", mutate)
+        result = _run(executor, {"Local.config": {"code": "existing", "enabled": True}, "simple_key": "v"})
+
+        assert result["shared_state_updates"] == {}
+
+    def test_deleted_key_reported(self) -> None:
+        """A key removed during execution is reported as a delete."""
+
+        async def mutate(message: Any, source_executor_ids: Any, state: Any, runner_context: Any) -> None:
+            state.delete("to_remove")
+            state.commit()
+
+        executor = _make_executor("test-exec", mutate)
+        result = _run(executor, {"to_remove": "value", "keep": "value"})
+
+        assert "to_remove" in result["shared_state_deletes"]
+        assert "keep" not in result["shared_state_deletes"]
+
+
+if __name__ == "__main__":
+    import pytest
+
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_workflow_client.py

@@ -0,0 +1,692 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for DurableWorkflowClient.
+
+Covers starting workflows, awaiting output (including error/timeout paths),
+parsing pending human-in-the-loop (HITL) requests from custom status, and
+sanitizing HITL responses before delivery.
+"""
+
+import json
+from dataclasses import dataclass
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import WorkflowEvent
+
+from agent_framework_durabletask import DurableWorkflowClient
+from agent_framework_durabletask._workflows.naming import workflow_orchestrator_name
+from agent_framework_durabletask._workflows.serialization import serialize_value, serialize_workflow_event
+
+
+@dataclass
+class _Receipt:
+    """Module-level dataclass so it is picklable by serialize_value."""
+
+    order_id: int
+    total: float
+
+
+@pytest.fixture
+def mock_client() -> Mock:
+    """Create a mock TaskHubGrpcClient."""
+    return Mock()
+
+
+@pytest.fixture
+def workflow_client(mock_client: Mock) -> DurableWorkflowClient:
+    """Create a DurableWorkflowClient wrapping the mock client."""
+    return DurableWorkflowClient(mock_client)
+
+
+class TestStartWorkflow:
+    """Test starting workflow orchestrations."""
+
+    def test_start_workflow_schedules_orchestrator(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """start_workflow schedules the per-workflow orchestration by name."""
+        mock_client.schedule_new_orchestration.return_value = "instance-1"
+
+        result = workflow_client.start_workflow(input="hello", workflow_name="orders")
+
+        assert result == "instance-1"
+        mock_client.schedule_new_orchestration.assert_called_once_with(
+            workflow_orchestrator_name("orders"), input="hello", instance_id=None
+        )
+
+    def test_start_workflow_passes_non_string_input_unchanged(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """Non-string payloads are forwarded as-is (no string coercion)."""
+        mock_client.schedule_new_orchestration.return_value = "instance-2"
+        payload = {"order_id": 42, "items": ["a", "b"]}
+
+        workflow_client.start_workflow(input=payload, workflow_name="orders")
+
+        _, kwargs = mock_client.schedule_new_orchestration.call_args
+        assert kwargs["input"] == payload
+
+    def test_start_workflow_strips_forged_subworkflow_envelope(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """Reserved sub-workflow envelope keys in client input are stripped at the boundary.
+
+        Only an internal child dispatch may carry these keys; if untrusted input could,
+        it would smuggle a payload onto the orchestrator's trusted (pickle) path.
+        """
+        mock_client.schedule_new_orchestration.return_value = "i"
+        forged = {"__subworkflow_input__": {"__pickled__": "evil", "__type__": "x"}, "real": 1}
+
+        workflow_client.start_workflow(input=forged, workflow_name="orders")
+
+        _, kwargs = mock_client.schedule_new_orchestration.call_args
+        assert kwargs["input"] == {"real": 1}
+        assert "__subworkflow_input__" not in kwargs["input"]
+
+    def test_start_workflow_forwards_instance_id(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """An explicit instance id is forwarded to the underlying client."""
+        mock_client.schedule_new_orchestration.return_value = "explicit-id"
+
+        workflow_client.start_workflow(input="x", workflow_name="orders", instance_id="explicit-id")
+
+        _, kwargs = mock_client.schedule_new_orchestration.call_args
+        assert kwargs["instance_id"] == "explicit-id"
+
+
+class TestWorkflowNameTargeting:
+    """Resolving the target workflow name from a default or per-call value."""
+
+    def test_uses_constructor_default(self, mock_client: Mock) -> None:
+        """A client default workflow name is used when none is passed per call."""
+        client = DurableWorkflowClient(mock_client, workflow_name="billing")
+        mock_client.schedule_new_orchestration.return_value = "i"
+
+        client.start_workflow(input="x")
+
+        mock_client.schedule_new_orchestration.assert_called_once_with(
+            workflow_orchestrator_name("billing"), input="x", instance_id=None
+        )
+
+    def test_per_call_overrides_default(self, mock_client: Mock) -> None:
+        """A per-call workflow name overrides the constructor default."""
+        client = DurableWorkflowClient(mock_client, workflow_name="billing")
+        mock_client.schedule_new_orchestration.return_value = "i"
+
+        client.start_workflow(input="x", workflow_name="orders")
+
+        mock_client.schedule_new_orchestration.assert_called_once_with(
+            workflow_orchestrator_name("orders"), input="x", instance_id=None
+        )
+
+    def test_raises_when_no_name_resolvable(self, workflow_client: DurableWorkflowClient) -> None:
+        """With no default and no per-call name, starting raises a clear error."""
+        with pytest.raises(ValueError, match="No workflow name"):
+            workflow_client.start_workflow(input="x")
+
+
+class TestOwnershipValidation:
+    """Opt-in validation that an instance belongs to the targeted workflow."""
+
+    def test_runtime_status_returns_none_for_foreign_instance(self, mock_client: Mock) -> None:
+        """A status query scoped to a workflow returns None for a foreign instance."""
+        client = DurableWorkflowClient(mock_client, workflow_name="orders")
+        state = Mock()
+        state.name = workflow_orchestrator_name("billing")  # different workflow
+        state.runtime_status.name = "RUNNING"
+        mock_client.get_orchestration_state.return_value = state
+
+        assert client.get_runtime_status("instance-1") is None
+
+    def test_runtime_status_returns_status_for_owned_instance(self, mock_client: Mock) -> None:
+        """A status query returns the status for an instance of the targeted workflow."""
+        client = DurableWorkflowClient(mock_client, workflow_name="orders")
+        state = Mock()
+        state.name = workflow_orchestrator_name("orders")
+        state.runtime_status.name = "RUNNING"
+        mock_client.get_orchestration_state.return_value = state
+
+        assert client.get_runtime_status("instance-1") == "RUNNING"
+
+    def test_pending_hitl_empty_for_foreign_instance(self, mock_client: Mock) -> None:
+        """Pending HITL is empty for an instance of a different workflow."""
+        client = DurableWorkflowClient(mock_client, workflow_name="orders")
+        state = Mock()
+        state.name = workflow_orchestrator_name("billing")
+        state.serialized_custom_status = json.dumps({"pending_requests": {"req-1": {"source_executor_id": "x"}}})
+        mock_client.get_orchestration_state.return_value = state
+
+        assert client.get_pending_hitl_requests("instance-1") == []
+
+    def test_send_hitl_rejects_foreign_instance(self, mock_client: Mock) -> None:
+        """Sending a HITL response to a foreign instance raises and does not deliver."""
+        client = DurableWorkflowClient(mock_client, workflow_name="orders")
+        state = Mock()
+        state.name = workflow_orchestrator_name("billing")
+        mock_client.get_orchestration_state.return_value = state
+
+        with pytest.raises(ValueError, match="does not belong"):
+            client.send_hitl_response("instance-1", "req-1", {"approved": True})
+
+        mock_client.raise_orchestration_event.assert_not_called()
+
+    def test_send_hitl_allows_owned_instance(self, mock_client: Mock) -> None:
+        """Sending a HITL response to an owned instance delivers the event."""
+        client = DurableWorkflowClient(mock_client, workflow_name="orders")
+        state = Mock()
+        state.name = workflow_orchestrator_name("orders")
+        mock_client.get_orchestration_state.return_value = state
+
+        client.send_hitl_response("instance-1", "req-1", {"approved": True})
+
+        mock_client.raise_orchestration_event.assert_called_once()
+
+
+class TestAwaitWorkflowOutput:
+    """Test awaiting workflow completion and output."""
+
+    def test_returns_deserialized_output_on_completion(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A COMPLETED workflow returns its deserialized output."""
+        metadata = Mock()
+        metadata.runtime_status.name = "COMPLETED"
+        metadata.serialized_output = json.dumps(["result"])
+        mock_client.wait_for_orchestration_completion.return_value = metadata
+
+        output = workflow_client.await_workflow_output("instance-1")
+
+        assert output == ["result"]
+
+    def test_returns_none_when_no_output(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """A COMPLETED workflow with no output returns None."""
+        metadata = Mock()
+        metadata.runtime_status.name = "COMPLETED"
+        metadata.serialized_output = None
+        mock_client.wait_for_orchestration_completion.return_value = metadata
+
+        assert workflow_client.await_workflow_output("instance-1") is None
+
+    def test_reconstructs_typed_outputs(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """Typed outputs encoded by the activity come back as objects, not marker dicts."""
+        receipt = _Receipt(order_id=7, total=19.99)
+        # The shared activity stores each yielded output via serialize_value(), so a
+        # typed object is persisted as a checkpoint-marker dict.
+        metadata = Mock()
+        metadata.runtime_status.name = "COMPLETED"
+        metadata.serialized_output = json.dumps([serialize_value(receipt)])
+        mock_client.wait_for_orchestration_completion.return_value = metadata
+
+        output = workflow_client.await_workflow_output("instance-1")
+
+        assert output == [receipt]
+        assert isinstance(output[0], _Receipt)
+
+    def test_raises_timeout_when_not_completed(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """A None metadata (no completion) raises TimeoutError."""
+        mock_client.wait_for_orchestration_completion.return_value = None
+
+        with pytest.raises(TimeoutError, match="did not complete"):
+            workflow_client.await_workflow_output("instance-1", timeout_seconds=5)
+
+    def test_raises_runtime_error_on_failed_status(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A non-COMPLETED status raises RuntimeError."""
+        metadata = Mock()
+        metadata.runtime_status.name = "FAILED"
+        metadata.serialized_output = "boom"
+        mock_client.wait_for_orchestration_completion.return_value = metadata
+
+        with pytest.raises(RuntimeError, match="status FAILED"):
+            workflow_client.await_workflow_output("instance-1")
+
+
+class TestGetRuntimeStatus:
+    """Test reading the workflow's runtime status."""
+
+    def test_returns_status_name(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """The runtime status name is returned when state is available."""
+        state = Mock()
+        state.runtime_status.name = "RUNNING"
+        mock_client.get_orchestration_state.return_value = state
+
+        assert workflow_client.get_runtime_status("instance-1") == "RUNNING"
+
+    def test_returns_none_when_no_state(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """No orchestration state yields None (status unknown)."""
+        mock_client.get_orchestration_state.return_value = None
+
+        assert workflow_client.get_runtime_status("instance-1") is None
+
+
+class TestGetPendingHitlRequests:
+    """Test parsing pending HITL requests from custom status."""
+
+    def _state_with_status(self, status: object) -> Mock:
+        state = Mock()
+        state.serialized_custom_status = json.dumps(status) if status is not None else None
+        return state
+
+    def test_returns_empty_when_no_state(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """No orchestration state yields an empty list."""
+        mock_client.get_orchestration_state.return_value = None
+
+        assert workflow_client.get_pending_hitl_requests("instance-1") == []
+
+    def test_returns_empty_when_status_blank(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """A blank custom status yields an empty list."""
+        state = Mock()
+        state.serialized_custom_status = ""
+        mock_client.get_orchestration_state.return_value = state
+
+        assert workflow_client.get_pending_hitl_requests("instance-1") == []
+
+    def test_returns_empty_on_invalid_json(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """Malformed custom status JSON yields an empty list."""
+        state = Mock()
+        state.serialized_custom_status = "{not-json"
+        mock_client.get_orchestration_state.return_value = state
+
+        assert workflow_client.get_pending_hitl_requests("instance-1") == []
+
+    def test_parses_pending_requests(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """Pending requests are normalized into the documented shape."""
+        status = {
+            "pending_requests": {
+                "req-1": {
+                    "request_id": "req-1",
+                    "source_executor_id": "approver",
+                    "data": {"prompt": "approve?"},
+                    "request_type": "ApprovalRequest",
+                    "response_type": "ApprovalResponse",
+                }
+            }
+        }
+        mock_client.get_orchestration_state.return_value = self._state_with_status(status)
+
+        requests = workflow_client.get_pending_hitl_requests("instance-1")
+
+        assert requests == [
+            {
+                "request_id": "req-1",
+                "source_executor_id": "approver",
+                "data": {"prompt": "approve?"},
+                "request_type": "ApprovalRequest",
+                "response_type": "ApprovalResponse",
+            }
+        ]
+
+    def test_falls_back_to_dict_key_for_request_id(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """When a request omits request_id, the dict key is used."""
+        status = {"pending_requests": {"req-key": {"source_executor_id": "x"}}}
+        mock_client.get_orchestration_state.return_value = self._state_with_status(status)
+
+        requests = workflow_client.get_pending_hitl_requests("instance-1")
+
+        assert requests[0]["request_id"] == "req-key"
+
+    def test_ignores_non_dict_entries(self, workflow_client: DurableWorkflowClient, mock_client: Mock) -> None:
+        """Non-dict request entries are skipped."""
+        status = {"pending_requests": {"req-1": "not-a-dict"}}
+        mock_client.get_orchestration_state.return_value = self._state_with_status(status)
+
+        assert workflow_client.get_pending_hitl_requests("instance-1") == []
+
+    def test_returns_empty_when_pending_not_dict(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A non-dict pending_requests field yields an empty list."""
+        status = {"pending_requests": ["unexpected"]}
+        mock_client.get_orchestration_state.return_value = self._state_with_status(status)
+
+        assert workflow_client.get_pending_hitl_requests("instance-1") == []
+
+
+class TestSendHitlResponse:
+    """Test delivering HITL responses."""
+
+    def test_raises_orchestration_event_with_request_id(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """The response is delivered as an external event named by request id."""
+        workflow_client.send_hitl_response("instance-1", "req-1", {"approved": True})
+
+        mock_client.raise_orchestration_event.assert_called_once()
+        _, kwargs = mock_client.raise_orchestration_event.call_args
+        assert kwargs["event_name"] == "req-1"
+        assert kwargs["data"] == {"approved": True}
+
+    def test_strips_pickle_markers_before_delivery(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A crafted pickle-marker payload is neutralized before reaching the worker.
+
+        The HITL response is sent to the worker which deserializes it, so a payload
+        carrying the checkpoint ``__pickled__`` marker must be stripped client-side
+        (regression guard for the strip_pickle_markers call in send_hitl_response).
+        """
+        malicious = {"__pickled__": "<crafted-base64-payload>", "approved": True}
+
+        workflow_client.send_hitl_response("instance-1", "req-1", malicious)
+
+        _, kwargs = mock_client.raise_orchestration_event.call_args
+        # The whole marker-bearing dict is neutralized (replaced with None) rather
+        # than forwarded, so it can never reach pickle.loads on the worker.
+        assert kwargs["data"] is None
+
+
+class TestStreamWorkflow:
+    """Test streaming typed workflow events by polling custom status."""
+
+    def _state(self, *, status: str, events: list[dict] | None = None) -> Mock:
+        state = Mock()
+        state.runtime_status.name = status
+        if events is None:
+            state.serialized_custom_status = None
+        else:
+            state.serialized_custom_status = json.dumps({"state": "running", "events": events})
+        return state
+
+    async def test_streams_events_in_order_until_terminal(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """Events accrue across polls and stream in order; streaming ends at a terminal state."""
+        # Each poll returns a growing accumulated event list, then a terminal status.
+        mock_client.get_orchestration_state.side_effect = [
+            self._state(status="RUNNING", events=[{"type": "executor_invoked", "executor_id": "a"}]),
+            self._state(
+                status="RUNNING",
+                events=[
+                    {"type": "executor_invoked", "executor_id": "a"},
+                    {"type": "executor_completed", "executor_id": "a"},
+                ],
+            ),
+            self._state(
+                status="COMPLETED",
+                events=[
+                    {"type": "executor_invoked", "executor_id": "a"},
+                    {"type": "executor_completed", "executor_id": "a"},
+                ],
+            ),
+        ]
+
+        seen = [event async for event in workflow_client.stream_workflow("instance-1", poll_interval_seconds=0)]
+
+        # Each accumulated event is yielded exactly once, in order, as a typed event.
+        assert all(isinstance(e, WorkflowEvent) for e in seen)
+        assert [e.type for e in seen] == ["executor_invoked", "executor_completed"]
+        assert [e.executor_id for e in seen] == ["a", "a"]
+
+    async def test_terminal_with_no_status_yields_nothing(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A workflow that is already terminal with no custom status streams no events."""
+        mock_client.get_orchestration_state.return_value = self._state(status="COMPLETED")
+
+        seen = [event async for event in workflow_client.stream_workflow("instance-1", poll_interval_seconds=0)]
+
+        assert seen == []
+
+    async def test_streams_typed_event_data_roundtrip(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """An output event's data is reconstructed into its original typed object."""
+        receipt = _Receipt(order_id=7, total=42.5)
+        serialized_event = serialize_workflow_event(WorkflowEvent("output", data=receipt, executor_id="processor"))
+        mock_client.get_orchestration_state.side_effect = [
+            self._state(status="RUNNING", events=[serialized_event]),
+            self._state(status="COMPLETED", events=[serialized_event]),
+        ]
+
+        seen = [event async for event in workflow_client.stream_workflow("instance-1", poll_interval_seconds=0)]
+
+        assert len(seen) == 1
+        assert isinstance(seen[0], WorkflowEvent)
+        assert seen[0].type == "output"
+        assert seen[0].executor_id == "processor"
+        assert seen[0].data == receipt
+
+
+class TestRunWorkflow:
+    """Test the async run_workflow convenience (start + optional wait)."""
+
+    async def test_waits_and_returns_output_by_default(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """By default run_workflow starts the workflow and returns its deserialized output."""
+        mock_client.schedule_new_orchestration.return_value = "instance-1"
+        metadata = Mock()
+        metadata.name = workflow_orchestrator_name("orders")
+        metadata.runtime_status.name = "COMPLETED"
+        metadata.serialized_output = json.dumps(["done"])
+        mock_client.wait_for_orchestration_completion.return_value = metadata
+
+        result = await workflow_client.run_workflow(input="hello", workflow_name="orders")
+
+        assert result == ["done"]
+        mock_client.schedule_new_orchestration.assert_called_once()
+        mock_client.wait_for_orchestration_completion.assert_called_once()
+
+    async def test_no_wait_returns_instance_id_without_awaiting(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """With wait=False, run_workflow returns the instance id and does not await completion."""
+        mock_client.schedule_new_orchestration.return_value = "instance-2"
+
+        result = await workflow_client.run_workflow(input="hello", workflow_name="orders", wait=False)
+
+        assert result == "instance-2"
+        mock_client.wait_for_orchestration_completion.assert_not_called()
+
+
+class TestSubworkflowHitl:
+    """Sub-workflow HITL: qualified request ids in/out (B2 single-surface addressing)."""
+
+    @staticmethod
+    def _states(mock_client: Mock, by_instance: dict[str, dict | None]) -> None:
+        """Wire get_orchestration_state to return a state per instance id.
+
+        Each value is the custom-status dict for that instance (or None for no
+        status). ``name`` is unset so ownership validation is skipped (these tests
+        construct the client without a workflow_name default).
+        """
+
+        def _get_state(instance_id: str) -> Mock | None:
+            if instance_id not in by_instance:
+                return None
+            status = by_instance[instance_id]
+            state = Mock()
+            state.serialized_custom_status = json.dumps(status) if status is not None else None
+            return state
+
+        mock_client.get_orchestration_state.side_effect = _get_state
+
+    def test_collects_nested_request_with_qualified_id(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A request pending in a child sub-workflow surfaces with an {executor}~{ordinal}~{id} id."""
+        self._states(
+            mock_client,
+            {
+                "parent": {"state": "running", "subworkflows": {"sub": ["child-1"]}},
+                "child-1": {
+                    "state": "waiting_for_human_input",
+                    "pending_requests": {"req-9": {"request_id": "req-9", "source_executor_id": "inner_node"}},
+                },
+            },
+        )
+
+        requests = workflow_client.get_pending_hitl_requests("parent")
+
+        assert len(requests) == 1
+        assert requests[0]["request_id"] == "sub~0~req-9"
+        assert requests[0]["source_executor_id"] == "inner_node"
+
+    def test_collects_parent_and_nested_requests_together(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """Top-level and nested pending requests are both returned (nested qualified)."""
+        self._states(
+            mock_client,
+            {
+                "parent": {
+                    "state": "waiting_for_human_input",
+                    "pending_requests": {"top-1": {"request_id": "top-1", "source_executor_id": "outer_node"}},
+                    "subworkflows": {"sub": ["child-1"]},
+                },
+                "child-1": {
+                    "state": "waiting_for_human_input",
+                    "pending_requests": {"inner-1": {"request_id": "inner-1", "source_executor_id": "inner_node"}},
+                },
+            },
+        )
+
+        ids = {r["request_id"] for r in workflow_client.get_pending_hitl_requests("parent")}
+
+        assert ids == {"top-1", "sub~0~inner-1"}
+
+    def test_collects_deeply_nested_request_with_full_path(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """Two levels of nesting accumulate a full {a}~{i}~{b}~{j}~{id} path."""
+        self._states(
+            mock_client,
+            {
+                "parent": {"state": "running", "subworkflows": {"mid": ["child-1"]}},
+                "child-1": {"state": "running", "subworkflows": {"leaf": ["child-2"]}},
+                "child-2": {
+                    "state": "waiting_for_human_input",
+                    "pending_requests": {"deep": {"request_id": "deep", "source_executor_id": "leaf_node"}},
+                },
+            },
+        )
+
+        requests = workflow_client.get_pending_hitl_requests("parent")
+
+        assert [r["request_id"] for r in requests] == ["mid~0~leaf~0~deep"]
+
+    def test_send_qualified_response_routes_to_child_instance(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A qualified id resolves to the owning child instance and bare request id."""
+        self._states(
+            mock_client,
+            {"parent": {"state": "running", "subworkflows": {"sub": ["child-1"]}}},
+        )
+
+        workflow_client.send_hitl_response("parent", "sub~0~req-9", {"approved": True})
+
+        mock_client.raise_orchestration_event.assert_called_once()
+        args, kwargs = mock_client.raise_orchestration_event.call_args
+        assert args[0] == "child-1"
+        assert kwargs["event_name"] == "req-9"
+        assert kwargs["data"] == {"approved": True}
+
+    def test_send_deeply_qualified_response_routes_to_leaf(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A two-level qualified id lands on the leaf child with the bare id."""
+        self._states(
+            mock_client,
+            {
+                "parent": {"state": "running", "subworkflows": {"mid": ["child-1"]}},
+                "child-1": {"state": "running", "subworkflows": {"leaf": ["child-2"]}},
+            },
+        )
+
+        workflow_client.send_hitl_response("parent", "mid~0~leaf~0~deep", {"ok": 1})
+
+        args, kwargs = mock_client.raise_orchestration_event.call_args
+        assert args[0] == "child-2"
+        assert kwargs["event_name"] == "deep"
+
+    def test_send_qualified_response_unknown_subworkflow_raises(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A qualified id for an inactive sub-workflow raises and delivers nothing."""
+        self._states(mock_client, {"parent": {"state": "running"}})  # no subworkflows map
+
+        with pytest.raises(ValueError, match="No active sub-workflow"):
+            workflow_client.send_hitl_response("parent", "sub~0~req-9", {"approved": True})
+
+        mock_client.raise_orchestration_event.assert_not_called()
+
+    def test_unqualified_response_still_targets_named_instance(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A plain (unqualified) request id targets the given instance directly."""
+        self._states(mock_client, {"parent": {"state": "waiting_for_human_input"}})
+
+        workflow_client.send_hitl_response("parent", "req-1", {"approved": True})
+
+        args, kwargs = mock_client.raise_orchestration_event.call_args
+        assert args[0] == "parent"
+        assert kwargs["event_name"] == "req-1"
+
+    def test_multiple_children_of_one_executor_stay_addressable(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """Two children dispatched by one node are qualified by ordinal, not collapsed."""
+        self._states(
+            mock_client,
+            {
+                "parent": {"state": "running", "subworkflows": {"sub": ["child-1", "child-2"]}},
+                "child-1": {
+                    "state": "waiting_for_human_input",
+                    "pending_requests": {"r1": {"request_id": "r1", "source_executor_id": "a"}},
+                },
+                "child-2": {
+                    "state": "waiting_for_human_input",
+                    "pending_requests": {"r2": {"request_id": "r2", "source_executor_id": "b"}},
+                },
+            },
+        )
+
+        ids = {r["request_id"] for r in workflow_client.get_pending_hitl_requests("parent")}
+        assert ids == {"sub~0~r1", "sub~1~r2"}
+
+        # The second child (ordinal 1) is reachable, not shadowed by the first.
+        workflow_client.send_hitl_response("parent", "sub~1~r2", {"ok": 1})
+        args, kwargs = mock_client.raise_orchestration_event.call_args
+        assert args[0] == "child-2"
+        assert kwargs["event_name"] == "r2"
+
+    def test_nested_leaf_request_id_with_double_colon_round_trips(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A functional sub-workflow's ``auto::N`` leaf id survives qualification and routing."""
+        self._states(
+            mock_client,
+            {
+                "parent": {"state": "running", "subworkflows": {"sub": ["child-1"]}},
+                "child-1": {
+                    "state": "waiting_for_human_input",
+                    "pending_requests": {"auto::0": {"request_id": "auto::0", "source_executor_id": "fn"}},
+                },
+            },
+        )
+
+        requests = workflow_client.get_pending_hitl_requests("parent")
+        assert [r["request_id"] for r in requests] == ["sub~0~auto::0"]
+
+        workflow_client.send_hitl_response("parent", "sub~0~auto::0", {"ok": 1})
+        args, kwargs = mock_client.raise_orchestration_event.call_args
+        assert args[0] == "child-1"
+        assert kwargs["event_name"] == "auto::0"
+
+    def test_top_level_auto_request_id_is_not_treated_as_nested(
+        self, workflow_client: DurableWorkflowClient, mock_client: Mock
+    ) -> None:
+        """A top-level ``auto::N`` id (contains ``::`` but no ``~``) routes to the instance itself."""
+        self._states(mock_client, {"parent": {"state": "waiting_for_human_input"}})
+
+        workflow_client.send_hitl_response("parent", "auto::0", {"approved": True})
+
+        args, kwargs = mock_client.raise_orchestration_event.call_args
+        assert args[0] == "parent"
+        assert kwargs["event_name"] == "auto::0"

python/packages/durabletask/tests/test_workflow_input_coercion.py

@@ -0,0 +1,135 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for workflow initial-input coercion (`_coerce_initial_input`).
+
+A durable workflow runs as a durable orchestration, so its initial payload
+arrives as plain JSON (no type markers). The shared engine reconstructs the
+start executor's declared input type from that JSON, mirroring in-process
+delivery. These tests pin that behavior across the relevant start-executor
+shapes.
+"""
+
+import json
+from dataclasses import dataclass
+from unittest.mock import Mock
+
+from agent_framework import AgentExecutor, Executor, WorkflowContext, handler
+from pydantic import BaseModel
+
+from agent_framework_durabletask._workflows.orchestrator import _coerce_initial_input
+
+
+@dataclass
+class _Submission:
+    content_id: str
+    title: str
+
+
+class _SubmissionModel(BaseModel):
+    content_id: str
+    title: str
+
+
+class _StrStart(Executor):
+    def __init__(self) -> None:
+        super().__init__(id="str_start")
+
+    @handler
+    async def run(self, message: str, ctx: WorkflowContext) -> None:  # pragma: no cover - never invoked
+        ...
+
+
+class _DataclassStart(Executor):
+    def __init__(self) -> None:
+        super().__init__(id="dc_start")
+
+    @handler
+    async def run(self, message: _Submission, ctx: WorkflowContext) -> None:  # pragma: no cover - never invoked
+        ...
+
+
+class _PydanticStart(Executor):
+    def __init__(self) -> None:
+        super().__init__(id="pyd_start")
+
+    @handler
+    async def run(self, message: _SubmissionModel, ctx: WorkflowContext) -> None:  # pragma: no cover - never invoked
+        ...
+
+
+def _workflow_with(executor: Executor | Mock) -> Mock:
+    workflow = Mock()
+    workflow.executors = {executor.id: executor}
+    workflow.start_executor_id = executor.id
+    return workflow
+
+
+class TestCoerceInitialInput:
+    """Test reconstruction of the initial workflow input by start-executor type."""
+
+    def test_str_start_passes_string_through(self) -> None:
+        workflow = _workflow_with(_StrStart())
+
+        assert _coerce_initial_input(workflow, "hello world") == "hello world"
+
+    def test_dataclass_start_reconstructs_from_dict(self) -> None:
+        workflow = _workflow_with(_DataclassStart())
+
+        result = _coerce_initial_input(workflow, {"content_id": "x", "title": "T"})
+
+        assert isinstance(result, _Submission)
+        assert result.content_id == "x"
+        assert result.title == "T"
+
+    def test_pydantic_start_reconstructs_from_dict(self) -> None:
+        workflow = _workflow_with(_PydanticStart())
+
+        result = _coerce_initial_input(workflow, {"content_id": "x", "title": "T"})
+
+        assert isinstance(result, _SubmissionModel)
+        assert result.content_id == "x"
+
+    def test_str_start_leaves_dict_unchanged(self) -> None:
+        """A str-typed start executor declares text; a dict is not coerced to str."""
+        workflow = _workflow_with(_StrStart())
+        payload = {"content_id": "x"}
+
+        assert _coerce_initial_input(workflow, payload) == payload
+
+    def test_agent_start_passes_string_through(self) -> None:
+        agent_executor = Mock(spec=AgentExecutor)
+        agent_executor.id = "agent"
+        workflow = _workflow_with(agent_executor)
+
+        assert _coerce_initial_input(workflow, "draft this email") == "draft this email"
+
+    def test_agent_start_stringifies_dict(self) -> None:
+        """Agents only consume text, so a structured payload is serialized to text."""
+        agent_executor = Mock(spec=AgentExecutor)
+        agent_executor.id = "agent"
+        workflow = _workflow_with(agent_executor)
+
+        result = _coerce_initial_input(workflow, {"email": "hi"})
+
+        assert result == json.dumps({"email": "hi"})
+
+    def test_missing_start_executor_passes_through(self) -> None:
+        workflow = Mock()
+        workflow.executors = {}
+        workflow.start_executor_id = "missing"
+        payload = {"a": 1}
+
+        assert _coerce_initial_input(workflow, payload) == payload
+
+    def test_pickle_marker_injection_is_neutralized(self) -> None:
+        """A crafted pickle-marker payload is stripped before reconstruction (no pickle RCE).
+
+        The initial workflow input is untrusted, so a dict carrying the checkpoint
+        ``__pickled__`` marker must be neutralized rather than flowing into
+        ``deserialize_value`` (which would ``pickle.loads`` it).
+        """
+        workflow = _workflow_with(_DataclassStart())
+        malicious = {"__pickled__": "<crafted-base64-payload>", "content_id": "x", "title": "T"}
+
+        # The marker-bearing dict is replaced with None, never unpickled or reconstructed.
+        assert _coerce_initial_input(workflow, malicious) is None

python/packages/durabletask/tests/test_workflow_naming.py

@@ -0,0 +1,172 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for the durable workflow naming helpers.
+
+These helpers derive the **stable** durable names a hosted workflow registers
+under. Stability matters: durable replay resumes an in-flight orchestration only
+if the orchestration name still resolves, so the round-trip
+(``workflow_orchestrator_name`` ↔ ``workflow_name_from_orchestrator``) and the
+validation rules (reject empty / malformed / auto-generated names) are the
+contract the multi-workflow hosting builds on.
+"""
+
+import uuid
+
+import pytest
+
+from agent_framework_durabletask import (
+    DURABLE_NAME_PREFIX,
+    is_auto_generated_workflow_name,
+    validate_executor_id,
+    validate_workflow_name,
+    workflow_name_from_orchestrator,
+    workflow_orchestrator_name,
+)
+from agent_framework_durabletask._workflows.naming import (
+    MAX_EXECUTOR_ID_LENGTH,
+    SUBWORKFLOW_REQUEST_SEPARATOR,
+    qualify_subworkflow_request_id,
+    split_subworkflow_request_id,
+)
+
+
+class TestWorkflowOrchestratorName:
+    """``workflow_orchestrator_name`` derives ``dafx-{name}`` for valid names."""
+
+    def test_prepends_prefix(self) -> None:
+        assert workflow_orchestrator_name("orders") == "dafx-orders"
+
+    def test_uses_shared_prefix_constant(self) -> None:
+        assert workflow_orchestrator_name("orders") == f"{DURABLE_NAME_PREFIX}orders"
+
+    @pytest.mark.parametrize("name", ["a", "Order_Processor", "spam-detection", "wf123"])
+    def test_accepts_valid_names(self, name: str) -> None:
+        assert workflow_orchestrator_name(name) == f"dafx-{name}"
+
+    @pytest.mark.parametrize("name", ["", "1abc", "has space", "bad/char", "emoji😀"])
+    def test_rejects_invalid_names(self, name: str) -> None:
+        with pytest.raises(ValueError):
+            workflow_orchestrator_name(name)
+
+
+class TestWorkflowNameRoundTrip:
+    """``workflow_name_from_orchestrator`` inverts ``workflow_orchestrator_name``."""
+
+    @pytest.mark.parametrize("name", ["orders", "Order_Processor", "spam-detection", "wf123"])
+    def test_round_trips(self, name: str) -> None:
+        orchestrator = workflow_orchestrator_name(name)
+        assert workflow_name_from_orchestrator(orchestrator) == name
+
+    def test_returns_none_without_prefix(self) -> None:
+        # A bare orchestration name (no dafx- prefix) is "not one of ours".
+        assert workflow_name_from_orchestrator("workflow_orchestrator") is None
+
+
+class TestValidateExecutorId:
+    """``validate_executor_id`` guards the durable-naming / nested-HITL contract."""
+
+    @pytest.mark.parametrize("executor_id", ["router", "agent_node", "reviewer-node", "a", "Step1"])
+    def test_accepts_ordinary_ids(self, executor_id: str) -> None:
+        validate_executor_id(executor_id)  # does not raise
+
+    def test_rejects_empty(self) -> None:
+        with pytest.raises(ValueError, match="non-empty"):
+            validate_executor_id("")
+
+    def test_rejects_id_containing_separator(self) -> None:
+        bad = f"a{SUBWORKFLOW_REQUEST_SEPARATOR}b"
+        with pytest.raises(ValueError, match="reserved sub-workflow request separator"):
+            validate_executor_id(bad)
+
+    def test_rejects_overly_long_id(self) -> None:
+        with pytest.raises(ValueError, match="too long"):
+            validate_executor_id("x" * (MAX_EXECUTOR_ID_LENGTH + 1))
+
+
+class TestSubworkflowRequestIdQualification:
+    """Round-trip of the ``{executor}~{ordinal}~{leaf}`` qualified-request-id scheme."""
+
+    def test_separator_is_url_safe_tilde(self) -> None:
+        # '~' is RFC 3986 unreserved and (unlike '::') never appears in core request ids.
+        assert SUBWORKFLOW_REQUEST_SEPARATOR == "~"
+
+    def test_qualify_then_split_round_trips(self) -> None:
+        qualified = qualify_subworkflow_request_id("sub", 2, "req-9")
+        assert qualified == "sub~2~req-9"
+        assert split_subworkflow_request_id(qualified) == ("sub", 2, "req-9")
+
+    def test_split_returns_none_for_bare_id(self) -> None:
+        assert split_subworkflow_request_id("req-9") is None
+
+    def test_split_preserves_double_colon_leaf(self) -> None:
+        # A functional workflow's ``auto::0`` leaf survives one peel as the remainder.
+        assert split_subworkflow_request_id("sub~0~auto::0") == ("sub", 0, "auto::0")
+
+    def test_split_treats_double_colon_only_id_as_bare(self) -> None:
+        # ``auto::0`` has no '~', so it is a bare leaf, not a nested hop.
+        assert split_subworkflow_request_id("auto::0") is None
+
+    def test_split_treats_non_integer_ordinal_as_bare(self) -> None:
+        # A value whose second segment is not an integer is not a structural hop.
+        assert split_subworkflow_request_id("a~b~c") is None
+
+    def test_nested_qualification_round_trips(self) -> None:
+        deep = qualify_subworkflow_request_id("mid", 0, qualify_subworkflow_request_id("leaf", 1, "deep"))
+        assert deep == "mid~0~leaf~1~deep"
+        hop = split_subworkflow_request_id(deep)
+        assert hop is not None
+        executor_id, ordinal, remainder = hop
+        assert (executor_id, ordinal) == ("mid", 0)
+        assert split_subworkflow_request_id(remainder) == ("leaf", 1, "deep")
+
+    def test_returns_none_for_prefix_only(self) -> None:
+        assert workflow_name_from_orchestrator(DURABLE_NAME_PREFIX) is None
+
+    def test_strips_only_leading_prefix(self) -> None:
+        # Reverse is meant for orchestration names; it strips just the prefix, so a
+        # scoped activity-style name returns the remainder verbatim.
+        assert workflow_name_from_orchestrator("dafx-orders-translator") == "orders-translator"
+
+
+class TestValidateWorkflowName:
+    """``validate_workflow_name`` rejects unstable / unsafe identities."""
+
+    @pytest.mark.parametrize("name", ["a", "A", "wf", "Order_Processor", "spam-detection", "x" * 63])
+    def test_accepts_valid(self, name: str) -> None:
+        validate_workflow_name(name)  # should not raise
+
+    def test_rejects_empty(self) -> None:
+        with pytest.raises(ValueError, match="non-empty"):
+            validate_workflow_name("")
+
+    @pytest.mark.parametrize("name", ["1abc", "-abc", "_abc", "has space", "bad/char", "a.b", "x" * 64])
+    def test_rejects_malformed(self, name: str) -> None:
+        with pytest.raises(ValueError, match="invalid"):
+            validate_workflow_name(name)
+
+    def test_rejects_auto_generated(self) -> None:
+        name = f"WorkflowBuilder-{uuid.uuid4()}"
+        with pytest.raises(ValueError, match="auto-generated"):
+            validate_workflow_name(name)
+
+
+class TestIsAutoGeneratedWorkflowName:
+    """``is_auto_generated_workflow_name`` detects WorkflowBuilder defaults."""
+
+    def test_detects_uuid_default(self) -> None:
+        assert is_auto_generated_workflow_name(f"WorkflowBuilder-{uuid.uuid4()}") is True
+
+    def test_detects_uppercase_hex_uuid(self) -> None:
+        assert is_auto_generated_workflow_name(f"WorkflowBuilder-{str(uuid.uuid4()).upper()}") is True
+
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "orders",
+            "WorkflowBuilder",
+            "WorkflowBuilder-not-a-uuid",
+            "MyWorkflowBuilder-3f2b1c0a-1234-5678-9abc-def012345678",
+        ],
+    )
+    def test_ignores_explicit_names(self, name: str) -> None:
+        assert is_auto_generated_workflow_name(name) is False

python/packages/durabletask/tests/test_workflow_registration.py

@@ -0,0 +1,189 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for plan_workflow_registration.
+
+Verifies the host-agnostic decision of which executors become durable entities
+(agent executors) versus durable activities (everything else), and that agent
+executors are carried whole so each host can register entities under the
+executor id the orchestrator dispatches to.
+"""
+
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import AgentExecutor, Executor, WorkflowExecutor
+
+from agent_framework_durabletask import (
+    WorkflowRegistrationPlan,
+    collect_hosted_workflows,
+    plan_workflow_registration,
+)
+
+
+def _agent_executor(executor_id: str, agent_name: str) -> Mock:
+    agent = Mock()
+    agent.name = agent_name
+    executor = Mock(spec=AgentExecutor)
+    executor.id = executor_id
+    executor.agent = agent
+    return executor
+
+
+def _activity_executor(executor_id: str) -> Mock:
+    executor = Mock(spec=Executor)
+    executor.id = executor_id
+    return executor
+
+
+def _subworkflow_executor(executor_id: str, inner_workflow: Mock) -> Mock:
+    executor = Mock(spec=WorkflowExecutor)
+    executor.id = executor_id
+    executor.workflow = inner_workflow
+    return executor
+
+
+def _workflow(name: str, executors: dict[str, Mock]) -> Mock:
+    workflow = Mock()
+    workflow.name = name
+    workflow.executors = executors
+    return workflow
+
+
+class TestPlanWorkflowRegistration:
+    """Test classification of workflow executors into durable primitives."""
+
+    def test_agent_executor_classified_as_entity(self) -> None:
+        """An AgentExecutor is carried whole in agent_executors."""
+        agent_exec = _agent_executor("reviewer-node", "Reviewer")
+        workflow = Mock()
+        workflow.executors = {"reviewer-node": agent_exec}
+
+        plan = plan_workflow_registration(workflow)
+
+        assert plan.agent_executors == [agent_exec]
+        assert plan.activity_executors == []
+
+    def test_non_agent_executor_classified_as_activity(self) -> None:
+        """A plain Executor is classified as an activity."""
+        activity_exec = _activity_executor("router-node")
+        workflow = Mock()
+        workflow.executors = {"router-node": activity_exec}
+
+        plan = plan_workflow_registration(workflow)
+
+        assert plan.agent_executors == []
+        assert plan.activity_executors == [activity_exec]
+
+    def test_mixed_executors_are_partitioned(self) -> None:
+        """Agent and non-agent executors are split into the correct buckets."""
+        agent_exec = _agent_executor("agent-node", "Agent")
+        activity_exec = _activity_executor("activity-node")
+        workflow = Mock()
+        workflow.executors = {"agent-node": agent_exec, "activity-node": activity_exec}
+
+        plan = plan_workflow_registration(workflow)
+
+        assert plan.agent_executors == [agent_exec]
+        assert plan.activity_executors == [activity_exec]
+
+    def test_agent_executor_id_is_preserved_when_distinct_from_name(self) -> None:
+        """The plan keeps the executor (and its id), not just the bare agent.
+
+        This is the core of the identity fix: dispatch targets the executor id,
+        so registration must be able to use the id even when it differs from
+        ``agent.name``.
+        """
+        agent_exec = _agent_executor("custom-executor-id", "ReusedAgentName")
+        workflow = Mock()
+        workflow.executors = {"custom-executor-id": agent_exec}
+
+        plan = plan_workflow_registration(workflow)
+
+        assert plan.agent_executors[0].id == "custom-executor-id"
+        assert plan.agent_executors[0].agent.name == "ReusedAgentName"
+
+    def test_returns_workflow_registration_plan(self) -> None:
+        """The return value is a WorkflowRegistrationPlan."""
+        workflow = Mock()
+        workflow.executors = {}
+
+        plan = plan_workflow_registration(workflow)
+
+        assert isinstance(plan, WorkflowRegistrationPlan)
+        assert plan.agent_executors == []
+        assert plan.activity_executors == []
+
+    def test_subworkflow_executor_classified_separately(self) -> None:
+        """A WorkflowExecutor goes to subworkflow_executors, not activities."""
+        inner = _workflow("inner", {})
+        sub_exec = _subworkflow_executor("sub-node", inner)
+        activity_exec = _activity_executor("router-node")
+        workflow = _workflow("outer", {"sub-node": sub_exec, "router-node": activity_exec})
+
+        plan = plan_workflow_registration(workflow)
+
+        assert plan.subworkflow_executors == [sub_exec]
+        assert plan.activity_executors == [activity_exec]
+        assert plan.agent_executors == []
+
+
+class TestCollectHostedWorkflows:
+    """Test the recursive walk over nested sub-workflows."""
+
+    def test_single_workflow_yields_itself(self) -> None:
+        workflow = _workflow("solo", {"node": _activity_executor("node")})
+
+        assert [w.name for w in collect_hosted_workflows(workflow)] == ["solo"]
+
+    def test_yields_nested_subworkflows_parent_first(self) -> None:
+        inner = _workflow("inner", {"leaf": _activity_executor("leaf")})
+        sub_exec = _subworkflow_executor("sub", inner)
+        outer = _workflow("outer", {"sub": sub_exec})
+
+        assert [w.name for w in collect_hosted_workflows(outer)] == ["outer", "inner"]
+
+    def test_dedupes_shared_subworkflow_by_name(self) -> None:
+        """A sub-workflow reused by two nodes is yielded once."""
+        inner = _workflow("shared", {"leaf": _activity_executor("leaf")})
+        sub_a = _subworkflow_executor("a", inner)
+        sub_b = _subworkflow_executor("b", inner)
+        outer = _workflow("outer", {"a": sub_a, "b": sub_b})
+
+        assert [w.name for w in collect_hosted_workflows(outer)] == ["outer", "shared"]
+
+    def test_walks_multiple_levels(self) -> None:
+        leaf = _workflow("leaf_wf", {"x": _activity_executor("x")})
+        mid = _workflow("mid_wf", {"l": _subworkflow_executor("l", leaf)})
+        top = _workflow("top_wf", {"m": _subworkflow_executor("m", mid)})
+
+        assert [w.name for w in collect_hosted_workflows(top)] == ["top_wf", "mid_wf", "leaf_wf"]
+
+    def test_rejects_two_different_workflows_sharing_a_name(self) -> None:
+        """Two different sub-workflow instances with the same name collide and raise."""
+        inner_a = _workflow("shared", {"x": _activity_executor("x")})
+        inner_b = _workflow("shared", {"y": _activity_executor("y")})  # different instance, same name
+        outer = _workflow("outer", {"a": _subworkflow_executor("a", inner_a), "b": _subworkflow_executor("b", inner_b)})
+
+        with pytest.raises(ValueError, match="collides"):
+            list(collect_hosted_workflows(outer))
+
+    def test_rejects_case_insensitive_name_collision(self) -> None:
+        """Two different instances whose names differ only by case collide and raise.
+
+        The route ownership guard compares the durable orchestration name
+        case-insensitively, so case-only name variants must be rejected here or one
+        workflow's routes could operate on the other's instances.
+        """
+        inner_a = _workflow("shared", {"x": _activity_executor("x")})
+        inner_b = _workflow("Shared", {"y": _activity_executor("y")})  # case-only difference
+        outer = _workflow("outer", {"a": _subworkflow_executor("a", inner_a), "b": _subworkflow_executor("b", inner_b)})
+
+        with pytest.raises(ValueError, match="collides"):
+            list(collect_hosted_workflows(outer))
+
+    def test_same_instance_reused_is_deduped_not_rejected(self) -> None:
+        """The same sub-workflow instance referenced by two nodes (fan-out) is yielded once."""
+        inner = _workflow("shared", {"x": _activity_executor("x")})
+        outer = _workflow("outer", {"a": _subworkflow_executor("a", inner), "b": _subworkflow_executor("b", inner)})
+
+        assert [w.name for w in collect_hosted_workflows(outer)] == ["outer", "shared"]

python/packages/durabletask/tests/test_workflow_routing.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for synchronous edge-condition evaluation on the durabletask host.
+
+Durable orchestrators run as generators and evaluate edge conditions
+synchronously. A condition that returns an awaitable cannot be evaluated in
+that context, so the edge is treated as *not matched* (not traversed).
+"""
+
+from agent_framework._workflows._edge import Edge  # pyright: ignore[reportPrivateImportUsage]
+
+from agent_framework_durabletask._workflows.orchestrator import _evaluate_edge_condition_sync
+
+
+class TestEvaluateEdgeConditionSync:
+    """Synchronous edge-condition evaluation semantics."""
+
+    def test_no_condition_traverses(self) -> None:
+        edge = Edge("a", "b")
+        assert _evaluate_edge_condition_sync(edge, {"x": 1}) is True
+
+    def test_sync_true_traverses(self) -> None:
+        edge = Edge("a", "b", condition=lambda m: m["ok"])
+        assert _evaluate_edge_condition_sync(edge, {"ok": True}) is True
+
+    def test_sync_false_does_not_traverse(self) -> None:
+        edge = Edge("a", "b", condition=lambda m: m["ok"])
+        assert _evaluate_edge_condition_sync(edge, {"ok": False}) is False
+
+    def test_async_condition_is_not_traversed(self) -> None:
+        # The durabletask host evaluates conditions synchronously; an async
+        # condition cannot be evaluated, so the edge is treated as not matched
+        # even though it would resolve True when awaited.
+        async def gate(_message: object) -> bool:
+            return True
+
+        edge = Edge("a", "b", condition=gate)
+        assert _evaluate_edge_condition_sync(edge, {"x": 1}) is False

python/packages/durabletask/tests/test_workflow_serialization.py

@@ -0,0 +1,451 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for workflow serialization helpers.
+
+``resolve_type`` is annotated ``type | None`` and its result flows into
+``reconstruct_to_type``, which calls ``issubclass``. A non-class attribute
+(function, module member, etc.) would raise ``TypeError`` there, so the
+resolver must only ever return actual classes.
+
+``deserialize_workflow_output`` reverses the per-output ``serialize_value``
+encoding the shared activity applies, so typed outputs are returned as the
+original objects rather than checkpoint-marker dicts.
+
+``serialize_value`` / ``deserialize_value`` are the internal codec; the
+round-trip, ``reconstruct_to_type``, and ``strip_pickle_markers`` suites below
+guard the type fidelity and the trust-boundary defense that neutralizes
+attacker-injected pickle/type markers before they can reach ``pickle.loads()``.
+"""
+
+import json
+from collections import OrderedDict
+from dataclasses import dataclass
+
+from agent_framework import (
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+    AgentResponse,
+    Message,
+    WorkflowEvent,
+)
+from pydantic import BaseModel
+
+from agent_framework_durabletask._workflows.serialization import (
+    SUBWORKFLOW_INPUT_KEY,
+    deserialize_value,
+    deserialize_workflow_event,
+    deserialize_workflow_output,
+    reconstruct_to_type,
+    resolve_type,
+    serialize_value,
+    serialize_workflow_event,
+    strip_pickle_markers,
+    strip_subworkflow_markers,
+)
+
+
+@dataclass
+class _Decision:
+    """Module-level dataclass so it is picklable by serialize_value."""
+
+    approved: bool
+    note: str
+
+
+class TestResolveType:
+    """Test that resolve_type only returns real classes."""
+
+    def test_resolves_a_real_class(self) -> None:
+        assert resolve_type("collections:OrderedDict") is OrderedDict
+
+    def test_returns_none_for_non_class_attribute(self) -> None:
+        # json.dumps is a function; if resolve_type returned it, issubclass()
+        # inside reconstruct_to_type() would raise TypeError at runtime.
+        assert resolve_type("json:dumps") is None
+
+    def test_returns_none_for_unknown_attribute(self) -> None:
+        assert resolve_type("json:DoesNotExist") is None
+
+    def test_returns_none_for_malformed_key(self) -> None:
+        assert resolve_type("not-a-valid-key") is None
+
+
+class TestDeserializeWorkflowOutput:
+    """Reconstruction of stored workflow outputs."""
+
+    def test_primitives_pass_through(self) -> None:
+        # Mirror the stored shape: a list of yielded outputs, JSON round-tripped.
+        stored = json.loads(json.dumps([serialize_value("hello"), serialize_value(42)]))
+
+        assert deserialize_workflow_output(stored) == ["hello", 42]
+
+    def test_typed_outputs_are_reconstructed(self) -> None:
+        # A typed object is stored as a checkpoint-marker dict; it must come back
+        # as the original object, not the marker dict.
+        decision = _Decision(approved=True, note="ok")
+        stored = json.loads(json.dumps([serialize_value(decision)]))
+
+        result = deserialize_workflow_output(stored)
+
+        assert result == [decision]
+        assert isinstance(result[0], _Decision)
+
+    def test_none_passes_through(self) -> None:
+        assert deserialize_workflow_output(None) is None
+
+
+@dataclass
+class _Approval:
+    """Module-level dataclass so it is picklable by serialize_value."""
+
+    reason: str
+
+
+def _roundtrip(event: WorkflowEvent) -> WorkflowEvent:
+    # Mirror the real path: serialize, JSON round-trip through the custom status,
+    # then reconstruct on the client.
+    return deserialize_workflow_event(json.loads(json.dumps(serialize_workflow_event(event))))
+
+
+class TestWorkflowEventRoundtrip:
+    """serialize_workflow_event / deserialize_workflow_event preserve event identity."""
+
+    def test_output_event_reconstructs_typed_data(self) -> None:
+        result = _roundtrip(WorkflowEvent("output", data=_Approval(reason="ok"), executor_id="writer"))
+
+        assert result.type == "output"
+        assert result.executor_id == "writer"
+        assert result.data == _Approval(reason="ok")
+        assert isinstance(result.data, _Approval)
+
+    def test_executor_completed_without_data_roundtrips_to_none(self) -> None:
+        result = _roundtrip(WorkflowEvent.executor_completed("reviewer"))
+
+        assert result.type == "executor_completed"
+        assert result.executor_id == "reviewer"
+        assert result.data is None
+
+    def test_iteration_tag_is_preserved(self) -> None:
+        # The orchestrator tags each event with its superstep before publishing.
+        serialized = serialize_workflow_event(WorkflowEvent.executor_invoked("writer"))
+        serialized["iteration"] = 3
+
+        result = deserialize_workflow_event(json.loads(json.dumps(serialized)))
+
+        assert result.type == "executor_invoked"
+        assert result.iteration == 3
+
+    def test_request_info_event_roundtrips(self) -> None:
+        event: WorkflowEvent = WorkflowEvent.request_info(
+            request_id="req-1",
+            source_executor_id="approver",
+            request_data=_Approval(reason="needs sign-off"),
+            response_type=bool,
+        )
+
+        result = _roundtrip(event)
+
+        assert result.type == "request_info"
+        assert result.request_id == "req-1"
+        assert result.source_executor_id == "approver"
+        assert result.response_type is bool
+        assert result.data == _Approval(reason="needs sign-off")
+
+
+# Module-level test types (must be importable for checkpoint encoding roundtrip).
+@dataclass
+class SampleData:
+    """Sample dataclass for testing checkpoint encoding roundtrip."""
+
+    name: str
+    value: int
+
+
+class SampleModel(BaseModel):
+    """Sample Pydantic model for testing checkpoint encoding roundtrip."""
+
+    title: str
+    count: int
+
+
+@dataclass
+class DataclassWithPydanticField:
+    """Dataclass containing a Pydantic model field for testing nested serialization."""
+
+    label: str
+    model: SampleModel
+
+
+class TestSerializationRoundtrip:
+    """``serialize_value`` / ``deserialize_value`` round-trip the typed objects used in workflows."""
+
+    def test_roundtrip_chat_message(self) -> None:
+        """Test Message survives encode → decode roundtrip."""
+        original = Message(role="user", contents=["Hello"])
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert isinstance(decoded, Message)
+        assert decoded.role == "user"
+
+    def test_roundtrip_agent_executor_request(self) -> None:
+        """Test AgentExecutorRequest with nested Messages roundtrips."""
+        original = AgentExecutorRequest(
+            messages=[Message(role="user", contents=["Hi"])],
+            should_respond=True,
+        )
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert isinstance(decoded, AgentExecutorRequest)
+        assert len(decoded.messages) == 1
+        assert isinstance(decoded.messages[0], Message)
+        assert decoded.should_respond is True
+
+    def test_roundtrip_agent_executor_response(self) -> None:
+        """Test AgentExecutorResponse with nested AgentResponse roundtrips."""
+        original = AgentExecutorResponse(
+            executor_id="test_exec",
+            agent_response=AgentResponse(messages=[Message(role="assistant", contents=["Reply"])]),
+            full_conversation=[Message(role="assistant", contents=["Reply"])],
+        )
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert isinstance(decoded, AgentExecutorResponse)
+        assert decoded.executor_id == "test_exec"
+        assert isinstance(decoded.agent_response, AgentResponse)
+
+    def test_roundtrip_dataclass(self) -> None:
+        """Test custom dataclass roundtrips."""
+        original = SampleData(name="test", value=42)
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert isinstance(decoded, SampleData)
+        assert decoded.name == "test"
+        assert decoded.value == 42
+
+    def test_roundtrip_pydantic_model(self) -> None:
+        """Test Pydantic model roundtrips."""
+        original = SampleModel(title="Hello", count=5)
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert isinstance(decoded, SampleModel)
+        assert decoded.title == "Hello"
+        assert decoded.count == 5
+
+    def test_roundtrip_primitives(self) -> None:
+        """Test primitives pass through unchanged."""
+        assert serialize_value(None) is None
+        assert serialize_value("hello") == "hello"
+        assert serialize_value(42) == 42
+        assert serialize_value(3.14) == 3.14
+        assert serialize_value(True) is True
+
+    def test_roundtrip_list_of_objects(self) -> None:
+        """Test list of typed objects roundtrips."""
+        original = [
+            Message(role="user", contents=["Q"]),
+            Message(role="assistant", contents=["A"]),
+        ]
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert isinstance(decoded, list)
+        assert len(decoded) == 2
+        assert all(isinstance(m, Message) for m in decoded)
+
+    def test_roundtrip_dict_of_objects(self) -> None:
+        """Test dict with typed values roundtrips (used for shared state)."""
+        original = {"count": 42, "msg": Message(role="user", contents=["Hi"])}
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert decoded["count"] == 42
+        assert isinstance(decoded["msg"], Message)
+
+    def test_roundtrip_dataclass_with_nested_pydantic(self) -> None:
+        """Test dataclass containing a Pydantic model field roundtrips correctly.
+
+        This covers the HITL pattern where AnalysisWithSubmission (dataclass)
+        contains a ContentAnalysisResult (Pydantic BaseModel) field.
+        """
+        original = DataclassWithPydanticField(label="test", model=SampleModel(title="Nested", count=99))
+        encoded = serialize_value(original)
+        decoded = deserialize_value(encoded)
+
+        assert isinstance(decoded, DataclassWithPydanticField)
+        assert decoded.label == "test"
+        assert isinstance(decoded.model, SampleModel)
+        assert decoded.model.title == "Nested"
+        assert decoded.model.count == 99
+
+
+class TestReconstructToType:
+    """Test suite for reconstruct_to_type function (used for HITL responses)."""
+
+    def test_none_returns_none(self) -> None:
+        """Test that None input returns None."""
+        assert reconstruct_to_type(None, str) is None
+
+    def test_already_correct_type(self) -> None:
+        """Test that values already of the correct type are returned as-is."""
+        assert reconstruct_to_type("hello", str) == "hello"
+        assert reconstruct_to_type(42, int) == 42
+
+    def test_non_dict_returns_original(self) -> None:
+        """Test that non-dict values are returned as-is."""
+        assert reconstruct_to_type("hello", int) == "hello"
+        assert reconstruct_to_type([1, 2], dict) == [1, 2]
+
+    def test_reconstruct_pydantic_model(self) -> None:
+        """Test reconstruction of Pydantic model from plain dict."""
+
+        class ApprovalResponse(BaseModel):
+            approved: bool
+            reason: str
+
+        data = {"approved": True, "reason": "Looks good"}
+        result = reconstruct_to_type(data, ApprovalResponse)
+
+        assert isinstance(result, ApprovalResponse)
+        assert result.approved is True
+        assert result.reason == "Looks good"
+
+    def test_reconstruct_dataclass(self) -> None:
+        """Test reconstruction of dataclass from plain dict."""
+
+        @dataclass
+        class Feedback:
+            score: int
+            comment: str
+
+        data = {"score": 5, "comment": "Great"}
+        result = reconstruct_to_type(data, Feedback)
+
+        assert isinstance(result, Feedback)
+        assert result.score == 5
+        assert result.comment == "Great"
+
+    def test_reconstruct_from_checkpoint_markers(self) -> None:
+        """Test that data with checkpoint markers is decoded via deserialize_value.
+
+        reconstruct_to_type is general-purpose and handles trusted checkpoint
+        data.  Untrusted HITL callers must call strip_pickle_markers() first.
+        """
+        original = SampleData(value=99, name="marker-test")
+        encoded = serialize_value(original)
+
+        result = reconstruct_to_type(encoded, SampleData)
+        assert isinstance(result, SampleData)
+        assert result.value == 99
+
+    def test_unrecognized_dict_returns_original(self) -> None:
+        """Test that unrecognized dicts are returned as-is."""
+
+        @dataclass
+        class Unrelated:
+            completely_different: str
+
+        data = {"some_key": "some_value"}
+        result = reconstruct_to_type(data, Unrelated)
+
+        assert result == data
+
+    def test_reconstruct_strips_injected_pickle_markers(self) -> None:
+        """End-to-end: strip_pickle_markers + reconstruct_to_type blocks attack.
+
+        This mirrors the real HITL flow where callers sanitize before reconstruction.
+        """
+        malicious = {"__pickled__": "gASVDgAAAAAAAACMBHRlc3SULg==", "__type__": "builtins:str"}
+        sanitized = strip_pickle_markers(malicious)
+        result = reconstruct_to_type(sanitized, str)
+        assert result is None
+
+
+class TestStripPickleMarkers:
+    """Security tests for strip_pickle_markers — the defence-in-depth layer
+    that prevents untrusted HTTP input from reaching pickle.loads()."""
+
+    def test_strips_top_level_pickle_marker(self) -> None:
+        """A dict containing __pickled__ must be replaced with None."""
+        data = {"__pickled__": "PAYLOAD", "__type__": "os:system"}
+        assert strip_pickle_markers(data) is None
+
+    def test_strips_top_level_type_marker_only(self) -> None:
+        """Even __type__ alone (without __pickled__) must be neutralised."""
+        data = {"__type__": "os:system", "other": "value"}
+        assert strip_pickle_markers(data) is None
+
+    def test_strips_nested_pickle_marker(self) -> None:
+        """Pickle markers nested inside a dict must be neutralised."""
+        data = {"safe": "value", "nested": {"__pickled__": "PAYLOAD", "__type__": "os:system"}}
+        result = strip_pickle_markers(data)
+        assert result == {"safe": "value", "nested": None}
+
+    def test_strips_pickle_marker_in_list(self) -> None:
+        """Pickle markers inside a list element must be neutralised."""
+        data = [{"__pickled__": "PAYLOAD"}, "safe"]
+        result = strip_pickle_markers(data)
+        assert result == [None, "safe"]
+
+    def test_strips_deeply_nested_marker(self) -> None:
+        """Deeply nested pickle markers must be neutralised."""
+        data = {"a": {"b": {"c": {"__pickled__": "deep"}}}}
+        result = strip_pickle_markers(data)
+        assert result == {"a": {"b": {"c": None}}}
+
+    def test_preserves_safe_dict(self) -> None:
+        """Dicts without pickle markers must be left untouched."""
+        data = {"approved": True, "reason": "Looks good"}
+        assert strip_pickle_markers(data) == data
+
+    def test_preserves_primitives(self) -> None:
+        """Primitive values must pass through unchanged."""
+        assert strip_pickle_markers("hello") == "hello"
+        assert strip_pickle_markers(42) == 42
+        assert strip_pickle_markers(None) is None
+        assert strip_pickle_markers(True) is True
+
+    def test_preserves_safe_list(self) -> None:
+        """Lists without pickle markers must be left untouched."""
+        data = [1, "two", {"key": "value"}]
+        assert strip_pickle_markers(data) == data
+
+    def test_mixed_safe_and_malicious(self) -> None:
+        """Only the malicious entries should be stripped; safe entries remain."""
+        data = {
+            "user_input": "hello",
+            "evil": {"__pickled__": "PAYLOAD", "__type__": "os:system"},
+            "count": 42,
+        }
+        result = strip_pickle_markers(data)
+        assert result == {"user_input": "hello", "evil": None, "count": 42}
+
+
+class TestStripSubworkflowMarkers:
+    """Boundary defence: a forged sub-workflow envelope in untrusted input is removed.
+
+    Only an internal child dispatch (post trust boundary) may carry the reserved
+    key; if untrusted client input could, it would be treated as a trusted
+    sub-orchestration payload and reach pickle.loads without sanitization.
+    """
+
+    def test_strips_input_key(self) -> None:
+        data = {SUBWORKFLOW_INPUT_KEY: {"__pickled__": "evil"}, "real": 1}
+        assert strip_subworkflow_markers(data) == {"real": 1}
+
+    def test_strips_full_forged_envelope(self) -> None:
+        data = {SUBWORKFLOW_INPUT_KEY: "x"}
+        assert strip_subworkflow_markers(data) == {}
+
+    def test_preserves_ordinary_dict(self) -> None:
+        data = {"order_id": 42, "items": ["a", "b"]}
+        assert strip_subworkflow_markers(data) == data
+
+    def test_preserves_non_dict(self) -> None:
+        assert strip_subworkflow_markers("hello") == "hello"
+        assert strip_subworkflow_markers([1, 2]) == [1, 2]
+        assert strip_subworkflow_markers(None) is None