chore: import upstream snapshot with attribution

python/packages/durabletask/AGENTS.md

@@ -0,0 +1,57 @@
+# Durable Task Package (agent-framework-durabletask)
+
+Durable execution support for long-running agent workflows using Azure Durable Functions.
+
+## Main Classes
+
+### Client Side
+
+- **`DurableAIAgentClient`** - Client for invoking durable agents
+- **`DurableAIAgent`** - Shim for creating durable agents
+
+### Worker Side
+
+- **`DurableAIAgentWorker`** - Worker that executes durable agent tasks
+- **`DurableAgentExecutor`** - Executes agent logic within durable context
+- **`AgentEntity`** - Durable entity for agent state management
+
+### State Management
+
+- **`DurableAgentState`** - State container for durable agents
+- **`DurableAgentSession`** - Session management for durable agents
+- **`DurableAIAgentOrchestrationContext`** - Orchestration context
+
+### Callbacks
+
+- **`AgentCallbackContext`** - Context for agent callbacks
+- **`AgentResponseCallbackProtocol`** - Protocol for response callbacks
+
+## Usage
+
+```python
+from agent_framework import Agent
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework_durabletask import DurableAIAgentClient, DurableAIAgentWorker
+from durabletask.client import TaskHubGrpcClient
+from durabletask.worker import TaskHubGrpcWorker
+
+# Client side
+dt_client = TaskHubGrpcClient(host_address="localhost:4001")
+agent_client = DurableAIAgentClient(dt_client)
+durable_agent = agent_client.get_agent("assistant")
+
+# Worker side
+dt_worker = TaskHubGrpcWorker(host_address="localhost:4001")
+agent_worker = DurableAIAgentWorker(dt_worker)
+
+# Create a chat client for the agent
+chat_client = OpenAIChatCompletionClient()
+my_agent = Agent(client=chat_client, name="assistant")
+agent_worker.add_agent(my_agent)
+```
+
+## Import Path
+
+```python
+from agent_framework_durabletask import DurableAIAgentClient, DurableAIAgentWorker
+```

python/packages/durabletask/LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Microsoft Corporation.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

python/packages/durabletask/README.md

@@ -0,0 +1,32 @@
+# Get Started with Microsoft Agent Framework Durable Task
+
+[![PyPI](https://img.shields.io/pypi/v/agent-framework-durabletask)](https://pypi.org/project/agent-framework-durabletask/)
+
+Please install this package via pip:
+
+```bash
+pip install agent-framework-durabletask --pre
+```
+
+## Durable Task Integration
+
+The durable task integration lets you host Microsoft Agent Framework agents using the [Durable Task](https://github.com/microsoft/durabletask-python) framework so they can persist state, replay conversation history, and recover from failures automatically.
+
+### Basic Usage Example
+
+```python
+from agent_framework import Agent
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework_durabletask import DurableAIAgentWorker
+from durabletask.worker import TaskHubGrpcWorker
+
+# Create the worker
+worker = TaskHubGrpcWorker(host_address="localhost:4001")
+agent_worker = DurableAIAgentWorker(worker)
+
+chat_client = OpenAIChatCompletionClient()
+my_agent = Agent(client=chat_client, name="assistant")
+agent_worker.add_agent(my_agent)
+```
+
+For more details, review the Python [README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) and the samples directory.

python/packages/durabletask/agent_framework_durabletask/__init__.py

@@ -0,0 +1,142 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Durable Task integration for Microsoft Agent Framework."""
+
+import importlib.metadata
+
+from ._async_bridge import run_agent_coroutine
+from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
+from ._client import DurableAIAgentClient
+from ._constants import (
+    DEFAULT_MAX_POLL_RETRIES,
+    DEFAULT_POLL_INTERVAL_SECONDS,
+    MIMETYPE_APPLICATION_JSON,
+    MIMETYPE_TEXT_PLAIN,
+    REQUEST_RESPONSE_FORMAT_JSON,
+    REQUEST_RESPONSE_FORMAT_TEXT,
+    THREAD_ID_FIELD,
+    THREAD_ID_HEADER,
+    WAIT_FOR_RESPONSE_FIELD,
+    WAIT_FOR_RESPONSE_HEADER,
+    ApiResponseFields,
+    ContentTypes,
+    DurableStateFields,
+)
+from ._durable_agent_state import (
+    DurableAgentState,
+    DurableAgentStateContent,
+    DurableAgentStateData,
+    DurableAgentStateDataContent,
+    DurableAgentStateEntry,
+    DurableAgentStateEntryJsonType,
+    DurableAgentStateErrorContent,
+    DurableAgentStateFunctionCallContent,
+    DurableAgentStateFunctionResultContent,
+    DurableAgentStateHostedFileContent,
+    DurableAgentStateHostedVectorStoreContent,
+    DurableAgentStateMessage,
+    DurableAgentStateRequest,
+    DurableAgentStateResponse,
+    DurableAgentStateTextContent,
+    DurableAgentStateTextReasoningContent,
+    DurableAgentStateUnknownContent,
+    DurableAgentStateUriContent,
+    DurableAgentStateUsage,
+    DurableAgentStateUsageContent,
+)
+from ._entities import AgentEntity, AgentEntityStateProviderMixin
+from ._executors import DurableAgentExecutor
+from ._models import AgentSessionId, DurableAgentSession, RunRequest
+from ._orchestration_context import DurableAIAgentOrchestrationContext
+from ._response_utils import ensure_response_format, load_agent_response
+from ._shim import DurableAIAgent
+from ._worker import DurableAIAgentWorker
+from ._workflows.activity import execute_workflow_activity
+from ._workflows.client import DurableWorkflowClient
+from ._workflows.context import WorkflowOrchestrationContext
+from ._workflows.dt_context import DurableTaskWorkflowContext
+from ._workflows.naming import (
+    DURABLE_NAME_PREFIX,
+    is_auto_generated_workflow_name,
+    validate_executor_id,
+    validate_workflow_name,
+    workflow_name_from_orchestrator,
+    workflow_orchestrator_name,
+)
+from ._workflows.orchestrator import run_workflow_orchestrator
+from ._workflows.registration import WorkflowRegistrationPlan, collect_hosted_workflows, plan_workflow_registration
+from ._workflows.runner_context import CapturingRunnerContext
+from ._workflows.serialization import deserialize_workflow_output
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"  # Fallback for development mode
+
+__all__ = [
+    "DEFAULT_MAX_POLL_RETRIES",
+    "DEFAULT_POLL_INTERVAL_SECONDS",
+    "DURABLE_NAME_PREFIX",
+    "MIMETYPE_APPLICATION_JSON",
+    "MIMETYPE_TEXT_PLAIN",
+    "REQUEST_RESPONSE_FORMAT_JSON",
+    "REQUEST_RESPONSE_FORMAT_TEXT",
+    "THREAD_ID_FIELD",
+    "THREAD_ID_HEADER",
+    "WAIT_FOR_RESPONSE_FIELD",
+    "WAIT_FOR_RESPONSE_HEADER",
+    "AgentCallbackContext",
+    "AgentEntity",
+    "AgentEntityStateProviderMixin",
+    "AgentResponseCallbackProtocol",
+    "AgentSessionId",
+    "ApiResponseFields",
+    "CapturingRunnerContext",
+    "ContentTypes",
+    "DurableAIAgent",
+    "DurableAIAgentClient",
+    "DurableAIAgentOrchestrationContext",
+    "DurableAIAgentWorker",
+    "DurableAgentExecutor",
+    "DurableAgentSession",
+    "DurableAgentState",
+    "DurableAgentStateContent",
+    "DurableAgentStateData",
+    "DurableAgentStateDataContent",
+    "DurableAgentStateEntry",
+    "DurableAgentStateEntryJsonType",
+    "DurableAgentStateErrorContent",
+    "DurableAgentStateFunctionCallContent",
+    "DurableAgentStateFunctionResultContent",
+    "DurableAgentStateHostedFileContent",
+    "DurableAgentStateHostedVectorStoreContent",
+    "DurableAgentStateMessage",
+    "DurableAgentStateRequest",
+    "DurableAgentStateResponse",
+    "DurableAgentStateTextContent",
+    "DurableAgentStateTextReasoningContent",
+    "DurableAgentStateUnknownContent",
+    "DurableAgentStateUriContent",
+    "DurableAgentStateUsage",
+    "DurableAgentStateUsageContent",
+    "DurableStateFields",
+    "DurableTaskWorkflowContext",
+    "DurableWorkflowClient",
+    "RunRequest",
+    "WorkflowOrchestrationContext",
+    "WorkflowRegistrationPlan",
+    "__version__",
+    "collect_hosted_workflows",
+    "deserialize_workflow_output",
+    "ensure_response_format",
+    "execute_workflow_activity",
+    "is_auto_generated_workflow_name",
+    "load_agent_response",
+    "plan_workflow_registration",
+    "run_agent_coroutine",
+    "run_workflow_orchestrator",
+    "validate_executor_id",
+    "validate_workflow_name",
+    "workflow_name_from_orchestrator",
+    "workflow_orchestrator_name",
+]

python/packages/durabletask/agent_framework_durabletask/_async_bridge.py

@@ -0,0 +1,89 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Persistent background event loop for running agent coroutines.
+
+Durable entity (and agent) handlers are invoked synchronously by the host on
+arbitrary worker threads. Agent clients and their async credentials create
+asyncio primitives (locks, connection pools, futures) that are bound to the
+event loop on which they are *first* used. Running a later invocation on a
+*different* event loop causes those primitives to await futures attached to a
+now-idle loop, which results in a silent, permanent hang.
+
+This module provides a single, process-wide persistent event loop running on a
+dedicated daemon thread. All agent coroutines are submitted to this loop via
+``run_coroutine_threadsafe`` so shared async resources remain valid across
+invocations regardless of which worker thread the host happens to use.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import threading
+from collections.abc import Coroutine
+from typing import Any, TypeVar
+
+_T = TypeVar("_T")
+
+_loop: asyncio.AbstractEventLoop | None = None
+_thread: threading.Thread | None = None
+_lock = threading.Lock()
+
+
+def _ensure_loop() -> asyncio.AbstractEventLoop:
+    """Return the shared persistent event loop, starting it on first use.
+
+    The loop is only reusable when it is open *and* its backing thread is still
+    alive. A loop whose thread has died (e.g. during interpreter shutdown) is not
+    reusable: ``run_coroutine_threadsafe`` would schedule onto a loop that will
+    never run again and ``future.result()`` would block forever. Such a loop is
+    replaced with a fresh loop + thread.
+    """
+    global _loop, _thread
+
+    loop, thread = _loop, _thread
+    if loop is not None and not loop.is_closed() and thread is not None and thread.is_alive():
+        return loop
+
+    with _lock:
+        loop, thread = _loop, _thread
+        if loop is not None and not loop.is_closed() and thread is not None and thread.is_alive():
+            return loop
+
+        # An existing loop whose thread has died is orphaned; close it best-effort
+        # before replacing it so it does not leak.
+        if loop is not None and not loop.is_closed():
+            with contextlib.suppress(Exception):
+                loop.close()
+
+        new_loop = asyncio.new_event_loop()
+
+        def _run() -> None:
+            asyncio.set_event_loop(new_loop)
+            new_loop.run_forever()
+
+        new_thread = threading.Thread(target=_run, name="dafx-agent-loop", daemon=True)
+        new_thread.start()
+
+        _loop = new_loop
+        _thread = new_thread
+        return new_loop
+
+
+def run_agent_coroutine(coro: Coroutine[Any, Any, _T]) -> _T:
+    """Run a coroutine on the shared persistent event loop and return its result.
+
+    The calling (worker) thread blocks until the coroutine completes. Because
+    every agent coroutine runs on the same loop, async resources created by
+    shared agent clients/credentials (locks, connection pools) remain bound to a
+    live loop across all invocations, preventing cross-loop hangs.
+
+    Args:
+        coro: The coroutine to execute.
+
+    Returns:
+        The coroutine's result.
+    """
+    loop = _ensure_loop()
+    future = asyncio.run_coroutine_threadsafe(coro, loop)
+    return future.result()

python/packages/durabletask/agent_framework_durabletask/_callbacks.py

@@ -0,0 +1,40 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Callback interfaces for Durable Agent executions.
+
+This module enables callers of AgentFunctionApp to supply streaming and final-response callbacks that are
+invoked during durable entity execution.
+"""
+
+from dataclasses import dataclass
+from typing import Protocol
+
+from agent_framework import AgentResponse, AgentResponseUpdate
+
+
+@dataclass(frozen=True)
+class AgentCallbackContext:
+    """Context supplied to callback invocations."""
+
+    agent_name: str
+    correlation_id: str
+    thread_id: str | None = None
+    request_message: str | None = None
+
+
+class AgentResponseCallbackProtocol(Protocol):
+    """Protocol describing the callbacks invoked during agent execution."""
+
+    async def on_streaming_response_update(
+        self,
+        update: AgentResponseUpdate,
+        context: AgentCallbackContext,
+    ) -> None:
+        """Handle a streaming response update emitted by the agent."""
+
+    async def on_agent_response(
+        self,
+        response: AgentResponse,
+        context: AgentCallbackContext,
+    ) -> None:
+        """Handle the final agent response."""

python/packages/durabletask/agent_framework_durabletask/_client.py

@@ -0,0 +1,92 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Client wrapper for Durable Task Agent Framework.
+
+This module provides the DurableAIAgentClient class for external clients to interact
+with durable agents via gRPC.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from agent_framework import AgentResponse
+from durabletask.client import TaskHubGrpcClient
+
+from ._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
+from ._executors import ClientAgentExecutor
+from ._shim import DurableAgentProvider, DurableAIAgent
+
+logger = logging.getLogger("agent_framework.durabletask")
+
+
+class DurableAIAgentClient(DurableAgentProvider[AgentResponse]):
+    """Client wrapper for interacting with durable agents externally.
+
+    This class wraps a durabletask TaskHubGrpcClient and provides a convenient
+    interface for retrieving and executing durable agents from external contexts.
+
+    Example:
+        ```python
+        from durabletask import TaskHubGrpcClient
+        from agent_framework.azure import DurableAIAgentClient
+
+        # Create the underlying client
+        client = TaskHubGrpcClient(host_address="localhost:4001")
+
+        # Wrap it with the agent client
+        agent_client = DurableAIAgentClient(client)
+
+        # Get an agent reference
+        agent = agent_client.get_agent("assistant")
+
+        # Run the agent (synchronous call that waits for completion)
+        response = agent.run("Hello, how are you?")
+        print(response.text)
+        ```
+    """
+
+    def __init__(
+        self,
+        client: TaskHubGrpcClient,
+        max_poll_retries: int = DEFAULT_MAX_POLL_RETRIES,
+        poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS,
+    ):
+        """Initialize the client wrapper.
+
+        Args:
+            client: The durabletask client instance to wrap
+            max_poll_retries: Maximum polling attempts when waiting for responses
+            poll_interval_seconds: Delay in seconds between polling attempts
+        """
+        self._client = client
+
+        # Validate and set polling parameters
+        self.max_poll_retries = max(1, max_poll_retries)
+        self.poll_interval_seconds = (
+            poll_interval_seconds if poll_interval_seconds > 0 else DEFAULT_POLL_INTERVAL_SECONDS
+        )
+
+        self._executor = ClientAgentExecutor(self._client, self.max_poll_retries, self.poll_interval_seconds)
+        logger.debug("[DurableAIAgentClient] Initialized with client type: %s", type(client).__name__)
+
+    def get_agent(self, agent_name: str) -> DurableAIAgent[AgentResponse]:
+        """Retrieve a DurableAIAgent shim for the specified agent.
+
+        This method returns a proxy object that can be used to execute the agent.
+        The actual agent must be registered on a worker with the same name.
+
+        Args:
+            agent_name: Name of the agent to retrieve (without the dafx- prefix)
+
+        Returns:
+            DurableAIAgent instance that can be used to run the agent
+
+        Note:
+            This method does not validate that the agent exists. Validation
+            will occur when the agent is executed. If the entity doesn't exist,
+            the execution will fail with an appropriate error.
+        """
+        logger.debug("[DurableAIAgentClient] Creating agent proxy for: %s", agent_name)
+
+        return DurableAIAgent(self._executor, agent_name)

python/packages/durabletask/agent_framework_durabletask/_constants.py

@@ -0,0 +1,130 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Constants for Azure Functions Agent Framework integration.
+
+This module contains:
+- Runtime configuration constants (polling, MIME types, headers)
+- JSON field name mappings for camelCase (JSON) ↔ snake_case (Python) serialization
+
+For serialization constants, use the DurableStateFields, ContentTypes, and EntryTypes classes
+to ensure consistent field naming between to_dict() and from_dict() methods.
+"""
+
+from typing import Final
+
+# Supported request/response formats and MIME types
+REQUEST_RESPONSE_FORMAT_JSON: str = "json"
+REQUEST_RESPONSE_FORMAT_TEXT: str = "text"
+MIMETYPE_APPLICATION_JSON: str = "application/json"
+MIMETYPE_TEXT_PLAIN: str = "text/plain"
+
+# Field and header names
+THREAD_ID_FIELD: str = "thread_id"
+THREAD_ID_HEADER: str = "x-ms-thread-id"
+WAIT_FOR_RESPONSE_FIELD: str = "wait_for_response"
+WAIT_FOR_RESPONSE_HEADER: str = "x-ms-wait-for-response"
+
+# Polling configuration
+DEFAULT_MAX_POLL_RETRIES: int = 30
+DEFAULT_POLL_INTERVAL_SECONDS: float = 1.0
+
+
+# =============================================================================
+# JSON Field Name Constants for Durable Agent State Serialization
+# =============================================================================
+# These constants ensure consistent camelCase field names in JSON serialization.
+# Use these in both to_dict() and from_dict() methods to prevent mismatches.
+
+# NOTE: Changing these constants is a breaking change and might require a schema version bump.
+
+
+class DurableStateFields:
+    """JSON field name constants for durable agent state serialization.
+
+    All field names are in camelCase to match the JSON schema.
+    Use these constants in both to_dict() and from_dict() methods.
+    """
+
+    # Schema-level fields
+    SCHEMA_VERSION: Final[str] = "schemaVersion"
+    DATA: Final[str] = "data"
+
+    # Entry discriminator
+    TYPE_DISCRIMINATOR: Final[str] = "$type"
+
+    # Internal field names
+    JSON_TYPE: Final[str] = "json_type"
+    TYPE_INTERNAL: Final[str] = "type"
+
+    # Common entry fields
+    CORRELATION_ID: Final[str] = "correlationId"
+    CREATED_AT: Final[str] = "createdAt"
+    MESSAGES: Final[str] = "messages"
+    EXTENSION_DATA: Final[str] = "extensionData"
+
+    # Request-specific fields
+    RESPONSE_TYPE: Final[str] = "responseType"
+    RESPONSE_SCHEMA: Final[str] = "responseSchema"
+    ORCHESTRATION_ID: Final[str] = "orchestrationId"
+
+    # Response-specific fields
+    USAGE: Final[str] = "usage"
+
+    # Message fields
+    ROLE: Final[str] = "role"
+    CONTENTS: Final[str] = "contents"
+    AUTHOR_NAME: Final[str] = "authorName"
+
+    # Content fields
+    TEXT: Final[str] = "text"
+    URI: Final[str] = "uri"
+    MEDIA_TYPE: Final[str] = "mediaType"
+    MESSAGE: Final[str] = "message"
+    ERROR_CODE: Final[str] = "errorCode"
+    DETAILS: Final[str] = "details"
+    CALL_ID: Final[str] = "callId"
+    NAME: Final[str] = "name"
+    ARGUMENTS: Final[str] = "arguments"
+    RESULT: Final[str] = "result"
+    FILE_ID: Final[str] = "fileId"
+    VECTOR_STORE_ID: Final[str] = "vectorStoreId"
+    CONTENT: Final[str] = "content"
+
+    # Usage fields (noqa: S105 - these are JSON field names, not passwords)
+    INPUT_TOKEN_COUNT: Final[str] = "inputTokenCount"  # noqa: S105
+    OUTPUT_TOKEN_COUNT: Final[str] = "outputTokenCount"  # noqa: S105
+    TOTAL_TOKEN_COUNT: Final[str] = "totalTokenCount"  # noqa: S105
+
+    # History field
+    CONVERSATION_HISTORY: Final[str] = "conversationHistory"
+
+
+class ContentTypes:
+    """Content type discriminator values for the $type field.
+
+    These values are used in the JSON $type field to identify content types.
+    """
+
+    TEXT: Final[str] = "text"
+    DATA: Final[str] = "data"
+    ERROR: Final[str] = "error"
+    FUNCTION_CALL: Final[str] = "functionCall"
+    FUNCTION_RESULT: Final[str] = "functionResult"
+    HOSTED_FILE: Final[str] = "hostedFile"
+    HOSTED_VECTOR_STORE: Final[str] = "hostedVectorStore"
+    REASONING: Final[str] = "reasoning"
+    URI: Final[str] = "uri"
+    USAGE: Final[str] = "usage"
+    UNKNOWN: Final[str] = "unknown"
+
+
+class ApiResponseFields:
+    """Field names for HTTP API responses (not part of persisted schema).
+
+    These are used in try_get_agent_response() for backward compatibility
+    with the HTTP API response format.
+    """
+
+    CONTENT: Final[str] = "content"
+    MESSAGE_COUNT: Final[str] = "message_count"
+    CORRELATION_ID: Final[str] = "correlationId"

python/packages/durabletask/agent_framework_durabletask/_entities.py

@@ -0,0 +1,353 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Durable Task entity implementations for Microsoft Agent Framework."""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from datetime import datetime, timezone
+from typing import Any, cast
+
+from agent_framework import (
+    AgentResponse,
+    AgentResponseUpdate,
+    Content,
+    Message,
+    ResponseStream,
+    SupportsAgentRun,
+)
+from durabletask.entities import DurableEntity
+
+from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
+from ._durable_agent_state import (
+    DurableAgentState,
+    DurableAgentStateEntry,
+    DurableAgentStateMessage,
+    DurableAgentStateRequest,
+    DurableAgentStateResponse,
+)
+from ._models import RunRequest
+
+logger = logging.getLogger("agent_framework.durabletask")
+
+
+class AgentEntityStateProviderMixin:
+    """Mixin implementing durable agent state caching + (de)serialization + persistence.
+
+    Concrete classes must implement:
+    - _get_state_dict(): fetch raw persisted state dict (default should be {})
+    - _set_state_dict(): persist raw state dict
+    - _get_thread_id_from_entity(): fetch the thread ID from the underlying context
+    """
+
+    _state_cache: DurableAgentState | None = None
+
+    def _get_state_dict(self) -> dict[str, Any]:
+        raise NotImplementedError
+
+    def _set_state_dict(self, state: dict[str, Any]) -> None:
+        raise NotImplementedError
+
+    def _get_thread_id_from_entity(self) -> str:
+        raise NotImplementedError
+
+    @property
+    def thread_id(self) -> str:
+        return self._get_thread_id_from_entity()
+
+    @property
+    def state(self) -> DurableAgentState:
+        if self._state_cache is None:
+            raw_state = self._get_state_dict()
+            self._state_cache = DurableAgentState.from_dict(raw_state) if raw_state else DurableAgentState()
+        return self._state_cache
+
+    @state.setter
+    def state(self, value: DurableAgentState) -> None:
+        self._state_cache = value
+        self.persist_state()
+
+    def persist_state(self) -> None:
+        """Persist the current state to the underlying storage provider."""
+        if self._state_cache is None:
+            self._state_cache = DurableAgentState()
+        self._set_state_dict(self._state_cache.to_dict())
+
+    def reset(self) -> None:
+        """Clear conversation history by resetting state to a fresh DurableAgentState."""
+        self._state_cache = DurableAgentState()
+        self.persist_state()
+        logger.debug("[AgentEntityStateProviderMixin.reset] State reset complete")
+
+
+class AgentEntity:
+    """Platform-agnostic agent execution logic.
+
+    This class encapsulates the core logic for executing an agent within a durable entity context.
+    """
+
+    agent: SupportsAgentRun
+    callback: AgentResponseCallbackProtocol | None
+
+    def __init__(
+        self,
+        agent: SupportsAgentRun,
+        callback: AgentResponseCallbackProtocol | None = None,
+        *,
+        state_provider: AgentEntityStateProviderMixin,
+    ) -> None:
+        self.agent = agent
+        self.callback = callback
+        self._state_provider = state_provider
+
+        logger.debug("[AgentEntity] Initialized with agent type: %s", type(agent).__name__)
+
+    @property
+    def state(self) -> DurableAgentState:
+        return self._state_provider.state
+
+    @state.setter
+    def state(self, value: DurableAgentState) -> None:
+        self._state_provider.state = value
+
+    def persist_state(self) -> None:
+        self._state_provider.persist_state()
+
+    def reset(self) -> None:
+        self._state_provider.reset()
+
+    def _is_error_response(self, entry: DurableAgentStateEntry) -> bool:
+        """Check if a conversation history entry is an error response."""
+        if isinstance(entry, DurableAgentStateResponse):
+            return entry.is_error
+        return False
+
+    async def run(
+        self,
+        request: RunRequest | dict[str, Any] | str,
+    ) -> AgentResponse:
+        """Execute the agent with a message."""
+        if isinstance(request, str):
+            run_request = RunRequest.from_json(request)
+        elif isinstance(request, dict):
+            run_request = RunRequest.from_dict(request)
+        else:
+            run_request = request
+
+        message = run_request.message
+        thread_id = self._state_provider.thread_id
+        correlation_id = run_request.correlation_id
+        if not thread_id:
+            raise ValueError("Entity State Provider must provide a thread_id")
+        options: dict[str, Any] = dict(run_request.options)
+        options.setdefault("response_format", run_request.response_format)
+        if not run_request.enable_tool_calls:
+            options.setdefault("tools", None)
+
+        logger.debug("[AgentEntity.run] Received ThreadId %s Message: %s", thread_id, run_request)
+
+        state_request = DurableAgentStateRequest.from_run_request(run_request)
+        self.state.data.conversation_history.append(state_request)
+
+        try:
+            chat_messages: list[Message] = [
+                replayable_message
+                for entry in self.state.data.conversation_history
+                if not self._is_error_response(entry)
+                for m in entry.messages
+                if (replayable_message := self._to_replayable_message(m)) is not None
+            ]
+
+            run_kwargs: dict[str, Any] = {"messages": chat_messages, "options": options}
+
+            agent_run_response: AgentResponse = await self._invoke_agent(
+                run_kwargs=run_kwargs,
+                correlation_id=correlation_id,
+                thread_id=thread_id,
+                request_message=message,
+            )
+
+            state_response = DurableAgentStateResponse.from_run_response(correlation_id, agent_run_response)
+            self.state.data.conversation_history.append(state_response)
+            self.persist_state()
+
+            return agent_run_response
+
+        except Exception as exc:
+            logger.exception("[AgentEntity.run] Agent execution failed.")
+
+            error_message = Message(
+                role="assistant", contents=[Content.from_error(message=str(exc), error_code=type(exc).__name__)]
+            )
+            error_response = AgentResponse(
+                messages=[error_message],
+                created_at=datetime.now(tz=timezone.utc).isoformat(),
+            )
+
+            error_state_response = DurableAgentStateResponse.from_run_response(correlation_id, error_response)
+            error_state_response.is_error = True
+            self.state.data.conversation_history.append(error_state_response)
+            self.persist_state()
+
+            return error_response
+
+    @staticmethod
+    def _to_replayable_message(message: DurableAgentStateMessage) -> Message | None:
+        """Convert persisted history into a message safe to replay into chat clients."""
+        chat_message = message.to_chat_message()
+        replayable_contents = [content for content in chat_message.contents if content.type != "reasoning"]
+        if not replayable_contents:
+            return None
+
+        return Message(
+            role=chat_message.role,
+            contents=replayable_contents,
+            author_name=chat_message.author_name,
+            additional_properties=chat_message.additional_properties,
+        )
+
+    async def _invoke_agent(
+        self,
+        run_kwargs: dict[str, Any],
+        correlation_id: str,
+        thread_id: str,
+        request_message: str,
+    ) -> AgentResponse:
+        """Execute the agent, preferring streaming when available."""
+        callback_context: AgentCallbackContext | None = None
+        if self.callback is not None:
+            callback_context = self._build_callback_context(
+                correlation_id=correlation_id,
+                thread_id=thread_id,
+                request_message=request_message,
+            )
+
+        run_callable = self.agent.run
+
+        # Try streaming first with run(stream=True)
+        try:
+            stream_candidate = run_callable(stream=True, **run_kwargs)
+            if inspect.isawaitable(stream_candidate):
+                stream_candidate = await stream_candidate
+
+            return await self._consume_stream(
+                stream=stream_candidate,
+                callback_context=callback_context,
+            )
+        except TypeError as type_error:
+            if "__aiter__" not in str(type_error) and "stream" not in str(type_error):
+                raise
+            logger.debug(
+                "run(stream=True) returned a non-async result; falling back to run(): %s",
+                type_error,
+            )
+        except Exception as stream_error:
+            logger.warning(
+                "run(stream=True) failed; falling back to run(): %s",
+                stream_error,
+                exc_info=True,
+            )
+        agent_run_response = run_callable(**run_kwargs)
+        if inspect.isawaitable(agent_run_response):
+            agent_run_response = await agent_run_response
+
+        if not isinstance(agent_run_response, AgentResponse):
+            raise TypeError(
+                f"Agent run() must return an AgentResponse instance; received {type(agent_run_response).__name__}"
+            )
+        await self._notify_final_response(agent_run_response, callback_context)
+        return agent_run_response
+
+    async def _consume_stream(
+        self,
+        stream: ResponseStream[AgentResponseUpdate, AgentResponse],
+        callback_context: AgentCallbackContext | None = None,
+    ) -> AgentResponse:
+        """Consume streaming responses and build the final AgentResponse."""
+        updates: list[AgentResponseUpdate] = []
+
+        async for update in stream:
+            updates.append(update)
+            await self._notify_stream_update(update, callback_context)
+
+        response = await stream.get_final_response()
+
+        await self._notify_final_response(response, callback_context)
+        return response
+
+    async def _notify_stream_update(
+        self,
+        update: AgentResponseUpdate,
+        context: AgentCallbackContext | None,
+    ) -> None:
+        """Invoke the streaming callback if one is registered."""
+        if self.callback is None or context is None:
+            return
+
+        try:
+            callback_result = self.callback.on_streaming_response_update(update, context)
+            if inspect.isawaitable(callback_result):
+                await callback_result
+        except Exception as exc:
+            logger.warning(
+                "[AgentEntity] Streaming callback raised an exception: %s",
+                exc,
+                exc_info=True,
+            )
+
+    async def _notify_final_response(
+        self,
+        response: AgentResponse,
+        context: AgentCallbackContext | None,
+    ) -> None:
+        """Invoke the final response callback if one is registered."""
+        if self.callback is None or context is None:
+            return
+
+        try:
+            callback_result = self.callback.on_agent_response(response, context)
+            if inspect.isawaitable(callback_result):
+                await callback_result
+        except Exception as exc:
+            logger.warning(
+                "[AgentEntity] Response callback raised an exception: %s",
+                exc,
+                exc_info=True,
+            )
+
+    def _build_callback_context(
+        self,
+        correlation_id: str,
+        thread_id: str,
+        request_message: str,
+    ) -> AgentCallbackContext:
+        """Create the callback context provided to consumers."""
+        agent_name = getattr(self.agent, "name", None) or type(self.agent).__name__
+        return AgentCallbackContext(
+            agent_name=agent_name,
+            correlation_id=correlation_id,
+            thread_id=thread_id,
+            request_message=request_message,
+        )
+
+
+class DurableTaskEntityStateProvider(DurableEntity, AgentEntityStateProviderMixin):
+    """DurableTask Durable Entity state provider for AgentEntity.
+
+    This class utilizes the Durable Entity context from `durabletask` package
+    to get and set the state of the agent entity.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def _get_state_dict(self) -> dict[str, Any]:
+        raw = self.get_state(dict, default={})
+        return cast(dict[str, Any], raw)
+
+    def _set_state_dict(self, state: dict[str, Any]) -> None:
+        self.set_state(state)
+
+    def _get_thread_id_from_entity(self) -> str:
+        return self.entity_context.entity_id.key

python/packages/durabletask/agent_framework_durabletask/_executors.py

@@ -0,0 +1,527 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Provider strategies for Durable Agent execution.
+
+These classes are internal execution strategies used by the DurableAIAgent shim.
+They are intentionally separate from the public client/orchestration APIs to keep
+only `get_agent` exposed to consumers. Executors implement the execution contract
+and are injected into the shim.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+from typing import Any, Generic, TypeVar
+
+from agent_framework import AgentResponse, AgentSession, Content, Message
+from durabletask.client import TaskHubGrpcClient
+from durabletask.entities import EntityInstanceId
+from durabletask.task import CompletableTask, CompositeTask, OrchestrationContext, Task
+from pydantic import BaseModel
+
+from ._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
+from ._durable_agent_state import DurableAgentState
+from ._models import AgentSessionId, DurableAgentSession, RunRequest
+from ._response_utils import ensure_response_format, load_agent_response
+
+logger = logging.getLogger("agent_framework.durabletask")
+
+# TypeVar for the task type returned by executors
+TaskT = TypeVar("TaskT")
+
+
+class DurableAgentTask(CompositeTask[AgentResponse], CompletableTask[AgentResponse]):
+    """A custom Task that wraps entity calls and provides typed AgentResponse results.
+
+    This task wraps the underlying entity call task and intercepts its completion
+    to convert the raw result into a typed AgentResponse object.
+
+    When yielded in an orchestration, this task returns an AgentResponse:
+        response: AgentResponse = yield durable_agent_task
+    """
+
+    def __init__(
+        self,
+        entity_task: CompletableTask[Any],
+        response_format: type[BaseModel] | None,
+        correlation_id: str,
+    ):
+        """Initialize the DurableAgentTask.
+
+        Args:
+            entity_task: The underlying entity call task
+            response_format: Optional Pydantic model for response parsing
+            correlation_id: Correlation ID for logging
+        """
+        self._response_format = response_format
+        self._correlation_id = correlation_id
+        super().__init__([entity_task])
+
+    def on_child_completed(self, task: Task[Any]) -> None:
+        """Handle completion of the underlying entity task.
+
+        Parameters
+        ----------
+        task : Task
+            The entity call task that just completed
+        """
+        if self.is_complete:
+            return
+
+        if task.is_failed:
+            # Propagate the failure - pass the original exception directly
+            self.fail("call_entity Task failed", task.get_exception())
+            return
+
+        # Task succeeded - transform the raw result
+        raw_result = task.get_result()
+        logger.debug(
+            "[DurableAgentTask] Converting raw result for correlation_id %s",
+            self._correlation_id,
+        )
+
+        try:
+            response = load_agent_response(raw_result)
+
+            if self._response_format is not None:
+                ensure_response_format(
+                    self._response_format,
+                    self._correlation_id,
+                    response,
+                )
+
+            # Set the typed AgentResponse as this task's result
+            self.complete(response)
+
+        except Exception as ex:
+            err_msg = "[DurableAgentTask] Failed to convert result for correlation_id: " + self._correlation_id
+            logger.exception(err_msg)
+            self.fail(err_msg, ex)
+
+
+class DurableAgentExecutor(ABC, Generic[TaskT]):
+    """Abstract base class for durable agent execution strategies.
+
+    Type Parameters:
+        TaskT: The task type returned by this executor
+    """
+
+    @abstractmethod
+    def run_durable_agent(
+        self,
+        agent_name: str,
+        run_request: RunRequest,
+        session: AgentSession | None = None,
+    ) -> TaskT:
+        """Execute the durable agent.
+
+        Returns:
+            TaskT: The task type specific to this executor implementation
+        """
+        raise NotImplementedError
+
+    def get_new_session(
+        self,
+        agent_name: str,
+        *,
+        session_id: str | None = None,
+        service_session_id: str | None = None,
+    ) -> DurableAgentSession:
+        """Create a new DurableAgentSession with random session ID."""
+        durable_session_id = self._create_session_id(agent_name)
+        return DurableAgentSession(
+            durable_session_id=durable_session_id,
+            session_id=session_id,
+            service_session_id=service_session_id,
+        )
+
+    def _create_session_id(
+        self,
+        agent_name: str,
+        session: AgentSession | None = None,
+    ) -> AgentSessionId:
+        """Create the AgentSessionId for the execution."""
+        if isinstance(session, DurableAgentSession) and session.durable_session_id is not None:
+            return session.durable_session_id
+        # Create new session ID - either no session provided or it's a regular AgentSession
+        key = self.generate_unique_id()
+        return AgentSessionId(name=agent_name, key=key)
+
+    def generate_unique_id(self) -> str:
+        """Generate a new Unique ID."""
+        return uuid.uuid4().hex
+
+    def get_run_request(
+        self,
+        message: str,
+        *,
+        options: dict[str, Any] | None = None,
+    ) -> RunRequest:
+        """Create a RunRequest from message and options."""
+        correlation_id = self.generate_unique_id()
+
+        # Create a copy to avoid modifying the caller's dict
+        opts = dict(options) if options else {}
+
+        # Extract and REMOVE known keys from options copy
+        response_format = opts.pop("response_format", None)
+        enable_tool_calls = opts.pop("enable_tool_calls", True)
+        wait_for_response = opts.pop("wait_for_response", True)
+
+        return RunRequest(
+            message=message,
+            response_format=response_format,
+            enable_tool_calls=enable_tool_calls,
+            wait_for_response=wait_for_response,
+            correlation_id=correlation_id,
+            options=opts,
+        )
+
+    def _create_acceptance_response(self, correlation_id: str) -> AgentResponse:
+        """Create an acceptance response for fire-and-forget mode.
+
+        Args:
+            correlation_id: Correlation ID for tracking the request
+
+        Returns:
+            AgentResponse: Acceptance response with correlation ID
+        """
+        acceptance_message = Message(
+            role="system",
+            contents=[
+                Content.from_text(
+                    f"Request accepted for processing (correlation_id: {correlation_id}). "
+                    f"Agent is executing in the background. "
+                    f"Retrieve response via your configured streaming or callback mechanism."
+                )
+            ],
+        )
+        return AgentResponse(
+            messages=[acceptance_message],
+            created_at=datetime.now(timezone.utc).isoformat(),
+        )
+
+
+class ClientAgentExecutor(DurableAgentExecutor[AgentResponse]):
+    """Execution strategy for external clients.
+
+    Note: Returns AgentResponse directly since the execution
+    is blocking until response is available via polling
+    as per the design of TaskHubGrpcClient.
+    """
+
+    def __init__(
+        self,
+        client: TaskHubGrpcClient,
+        max_poll_retries: int = DEFAULT_MAX_POLL_RETRIES,
+        poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS,
+    ):
+        self._client = client
+        self.max_poll_retries = max_poll_retries
+        self.poll_interval_seconds = poll_interval_seconds
+
+    def run_durable_agent(
+        self,
+        agent_name: str,
+        run_request: RunRequest,
+        session: AgentSession | None = None,
+    ) -> AgentResponse:
+        """Execute the agent via the durabletask client.
+
+        Signals the agent entity with a message request, then polls the entity
+        state to retrieve the response once processing is complete.
+
+        Note: This is a blocking/synchronous operation (in line with how
+        TaskHubGrpcClient works) that polls until a response is available or
+        timeout occurs.
+
+        Args:
+            agent_name: Name of the agent to execute
+            run_request: The run request containing message and optional response format
+            session: Optional conversation session (creates new if not provided)
+
+        Returns:
+            AgentResponse: The agent's response after execution completes, or an immediate
+                            acknowledgement if wait_for_response is False
+        """
+        # Signal the entity with the request
+        entity_id = self._signal_agent_entity(agent_name, run_request, session)
+
+        # If fire-and-forget mode, return immediately without polling
+        if not run_request.wait_for_response:
+            logger.info(
+                "[ClientAgentExecutor] Fire-and-forget mode: request signaled (correlation: %s)",
+                run_request.correlation_id,
+            )
+            return self._create_acceptance_response(run_request.correlation_id)
+
+        # Poll for the response
+        agent_response = self._poll_for_agent_response(entity_id, run_request.correlation_id)
+
+        # Handle and return the result
+        return self._handle_agent_response(agent_response, run_request.response_format, run_request.correlation_id)
+
+    def _signal_agent_entity(
+        self,
+        agent_name: str,
+        run_request: RunRequest,
+        session: AgentSession | None,
+    ) -> EntityInstanceId:
+        """Signal the agent entity with a run request.
+
+        Args:
+            agent_name: Name of the agent to execute
+            run_request: The run request containing message and optional response format
+            session: Optional conversation session
+
+        Returns:
+            entity_id
+        """
+        # Get or create session ID
+        session_id = self._create_session_id(agent_name, session)
+
+        # Create the entity ID
+        entity_id = EntityInstanceId(
+            entity=session_id.entity_name,
+            key=session_id.key,
+        )
+
+        logger.debug(
+            "[ClientAgentExecutor] Signaling entity '%s' (session: %s, correlation: %s)",
+            agent_name,
+            session_id,
+            run_request.correlation_id,
+        )
+
+        self._client.signal_entity(entity_id, "run", run_request.to_dict())
+        return entity_id
+
+    def _poll_for_agent_response(
+        self,
+        entity_id: EntityInstanceId,
+        correlation_id: str,
+    ) -> AgentResponse | None:
+        """Poll the entity for a response with retries.
+
+        Args:
+            entity_id: Entity instance identifier
+            correlation_id: Correlation ID to track the request
+
+        Returns:
+            The agent response if found, None if timeout occurs
+        """
+        agent_response = None
+
+        for attempt in range(1, self.max_poll_retries + 1):
+            # Initial sleep is intentional - give the entity time to process before first poll
+            time.sleep(self.poll_interval_seconds)
+
+            agent_response = self._poll_entity_for_response(entity_id, correlation_id)
+            if agent_response is not None:
+                logger.info(
+                    "[ClientAgentExecutor] Found response (attempt %d/%d, correlation: %s)",
+                    attempt,
+                    self.max_poll_retries,
+                    correlation_id,
+                )
+                break
+
+            logger.debug(
+                "[ClientAgentExecutor] Response not ready (attempt %d/%d)",
+                attempt,
+                self.max_poll_retries,
+            )
+
+        return agent_response
+
+    def _handle_agent_response(
+        self,
+        agent_response: AgentResponse | None,
+        response_format: type[BaseModel] | None,
+        correlation_id: str,
+    ) -> AgentResponse:
+        """Handle the agent response or create an error response.
+
+        Args:
+            agent_response: The response from polling, or None if timeout
+            response_format: Optional response format for validation
+            correlation_id: Correlation ID for logging
+
+        Returns:
+            AgentResponse with either the agent's response or an error message
+        """
+        if agent_response is not None:
+            try:
+                # Validate response format if specified
+                if response_format is not None:
+                    ensure_response_format(
+                        response_format,
+                        correlation_id,
+                        agent_response,
+                    )
+
+                return agent_response
+
+            except Exception as e:
+                logger.exception(
+                    "[ClientAgentExecutor] Error converting response for correlation: %s",
+                    correlation_id,
+                )
+                error_message = Message(
+                    role="system",
+                    contents=[
+                        Content.from_error(
+                            message=f"Error processing agent response: {e}",
+                            error_code="response_processing_error",
+                        )
+                    ],
+                )
+        else:
+            logger.warning(
+                "[ClientAgentExecutor] Timeout after %d attempts (correlation: %s)",
+                self.max_poll_retries,
+                correlation_id,
+            )
+            error_message = Message(
+                role="system",
+                contents=[
+                    Content.from_error(
+                        message=f"Timeout waiting for agent response after {self.max_poll_retries} attempts",
+                        error_code="response_timeout",
+                    )
+                ],
+            )
+
+        return AgentResponse(
+            messages=[error_message],
+            created_at=datetime.now(timezone.utc).isoformat(),
+        )
+
+    def _poll_entity_for_response(
+        self,
+        entity_id: EntityInstanceId,
+        correlation_id: str,
+    ) -> AgentResponse | None:
+        """Poll the entity state for a response matching the correlation ID.
+
+        Args:
+            entity_id: Entity instance identifier
+            correlation_id: Correlation ID to search for
+
+        Returns:
+            Response AgentResponse, None otherwise
+        """
+        try:
+            entity_metadata = self._client.get_entity(entity_id, include_state=True)
+
+            if entity_metadata is None:
+                return None
+
+            state_json = entity_metadata.get_state()
+            if not state_json:
+                return None
+
+            state = DurableAgentState.from_json(state_json)
+
+            # Use the helper method to get response by correlation ID
+            return state.try_get_agent_response(correlation_id)
+
+        except Exception as e:
+            logger.warning(
+                "[ClientAgentExecutor] Error reading entity state: %s",
+                e,
+            )
+            return None
+
+
+class OrchestrationAgentExecutor(DurableAgentExecutor[DurableAgentTask]):
+    """Execution strategy for orchestrations (sync/yield)."""
+
+    def __init__(self, context: OrchestrationContext):
+        self._context = context
+        logger.debug("[OrchestrationAgentExecutor] Initialized")
+
+    def generate_unique_id(self) -> str:
+        """Create a new UUID that is safe for replay within an orchestration or operation."""
+        return self._context.new_uuid()
+
+    def get_run_request(
+        self,
+        message: str,
+        *,
+        options: dict[str, Any] | None = None,
+    ) -> RunRequest:
+        """Get the current run request from the orchestration context.
+
+        Returns:
+            RunRequest: The current run request
+        """
+        request = super().get_run_request(
+            message,
+            options=options,
+        )
+        request.orchestration_id = self._context.instance_id
+        return request
+
+    def run_durable_agent(
+        self,
+        agent_name: str,
+        run_request: RunRequest,
+        session: AgentSession | None = None,
+    ) -> DurableAgentTask:
+        """Execute the agent via orchestration context.
+
+        Calls the agent entity and returns a DurableAgentTask that can be yielded
+        in orchestrations to wait for the entity's response.
+
+        Args:
+            agent_name: Name of the agent to execute
+            run_request: The run request containing message and optional response format
+            session: Optional conversation session (creates new if not provided)
+
+        Returns:
+            DurableAgentTask: A task wrapping the entity call that yields AgentResponse
+        """
+        # Resolve session
+        session_id = self._create_session_id(agent_name, session)
+
+        # Create the entity ID
+        entity_id = EntityInstanceId(
+            entity=session_id.entity_name,
+            key=session_id.key,
+        )
+
+        logger.debug(
+            "[OrchestrationAgentExecutor] correlation_id: %s entity_id: %s session_id: %s",
+            run_request.correlation_id,
+            entity_id,
+            session_id,
+        )
+
+        # Branch based on wait_for_response
+        if not run_request.wait_for_response:
+            # Fire-and-forget mode: signal entity and return pre-completed task
+            logger.info(
+                "[OrchestrationAgentExecutor] Fire-and-forget mode: signaling entity (correlation: %s)",
+                run_request.correlation_id,
+            )
+            self._context.signal_entity(entity_id, "run", run_request.to_dict())
+
+            # Create a pre-completed task with acceptance response
+            acceptance_response = self._create_acceptance_response(run_request.correlation_id)
+            entity_task: CompletableTask[AgentResponse] = CompletableTask()
+            entity_task.complete(acceptance_response)
+        else:
+            # Blocking mode: call entity and wait for response
+            entity_task = self._context.call_entity(entity_id, "run", run_request.to_dict())
+
+        # Wrap in DurableAgentTask for response transformation
+        return DurableAgentTask(
+            entity_task=entity_task,
+            response_format=run_request.response_format,
+            correlation_id=run_request.correlation_id,
+        )

python/packages/durabletask/agent_framework_durabletask/_models.py

@@ -0,0 +1,334 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Data models for Durable Agent Framework.
+
+This module defines the request and response models used by the framework.
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from importlib import import_module
+from typing import TYPE_CHECKING, Any, cast
+
+from agent_framework import AgentSession
+
+from ._constants import REQUEST_RESPONSE_FORMAT_TEXT
+
+if TYPE_CHECKING:  # pragma: no cover - type checking imports only
+    from pydantic import BaseModel
+
+_PydanticBaseModel: type[BaseModel] | None
+
+try:
+    from pydantic import BaseModel as _RuntimeBaseModel
+except ImportError:  # pragma: no cover - optional dependency
+    _PydanticBaseModel = None
+else:
+    _PydanticBaseModel = _RuntimeBaseModel
+
+
+def serialize_response_format(response_format: type[BaseModel] | None) -> Any:
+    """Serialize response format for transport across durable function boundaries."""
+    if response_format is None:
+        return None
+
+    if _PydanticBaseModel is None:
+        raise RuntimeError("pydantic is required to use structured response formats")
+
+    if not inspect.isclass(response_format) or not issubclass(response_format, _PydanticBaseModel):
+        raise TypeError("response_format must be a Pydantic BaseModel type")
+
+    return {
+        "__response_schema_type__": "pydantic_model",
+        "module": response_format.__module__,
+        "qualname": response_format.__qualname__,
+    }
+
+
+def _deserialize_response_format(response_format: Any) -> type[BaseModel] | None:
+    """Deserialize response format back into actionable type if possible."""
+    if response_format is None:
+        return None
+
+    if (
+        _PydanticBaseModel is not None
+        and inspect.isclass(response_format)
+        and issubclass(response_format, _PydanticBaseModel)
+    ):
+        return response_format
+
+    if not isinstance(response_format, dict):
+        return None
+
+    response_dict = cast(dict[str, Any], response_format)
+
+    if response_dict.get("__response_schema_type__") != "pydantic_model":
+        return None
+
+    module_name = response_dict.get("module")
+    qualname = response_dict.get("qualname")
+    if not module_name or not qualname:
+        return None
+
+    try:
+        module = import_module(module_name)
+    except ImportError:  # pragma: no cover - user provided module missing
+        return None
+
+    attr: Any = module
+    for part in qualname.split("."):
+        try:
+            attr = getattr(attr, part)
+        except AttributeError:  # pragma: no cover - invalid qualname
+            return None
+
+    if _PydanticBaseModel is not None and inspect.isclass(attr) and issubclass(attr, _PydanticBaseModel):
+        return attr
+
+    return None
+
+
+@dataclass
+class RunRequest:
+    """Represents a request to run an agent with a specific message and configuration.
+
+    Attributes:
+        message: The message to send to the agent
+        request_response_format: The desired response format (e.g., "text" or "json")
+        role: The role of the message sender (user, system, or assistant)
+        response_format: Optional Pydantic BaseModel type describing the structured response format
+        enable_tool_calls: Whether to enable tool calls for this request
+        wait_for_response: If True (default), caller will wait for agent response. If False,
+                          returns immediately after signaling (fire-and-forget mode)
+        correlation_id: Correlation ID for tracking the response to this specific request
+        created_at: Optional timestamp when the request was created
+        orchestration_id: Optional ID of the orchestration that initiated this request
+        options: Optional options dictionary forwarded to the agent
+    """
+
+    message: str
+    request_response_format: str
+    correlation_id: str
+    role: str = "user"
+    response_format: type[BaseModel] | None = None
+    enable_tool_calls: bool = True
+    wait_for_response: bool = True
+    created_at: datetime | None = None
+    orchestration_id: str | None = None
+    options: dict[str, Any] = field(default_factory=lambda: {})
+
+    def __init__(
+        self,
+        message: str,
+        correlation_id: str,
+        request_response_format: str = REQUEST_RESPONSE_FORMAT_TEXT,
+        role: str | None = "user",
+        response_format: type[BaseModel] | None = None,
+        enable_tool_calls: bool = True,
+        wait_for_response: bool = True,
+        created_at: datetime | None = None,
+        orchestration_id: str | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> None:
+        self.message = message
+        self.correlation_id = correlation_id
+        self.role = self.coerce_role(role)
+        self.response_format = response_format
+        self.request_response_format = request_response_format
+        self.enable_tool_calls = enable_tool_calls
+        self.wait_for_response = wait_for_response
+        self.created_at = created_at if created_at is not None else datetime.now(tz=timezone.utc)
+        self.orchestration_id = orchestration_id
+        self.options = options if options is not None else {}
+
+    @staticmethod
+    def coerce_role(value: str | None) -> str:
+        """Normalize various role representations into a role string."""
+        if isinstance(value, str):
+            normalized = value.strip()
+            if not normalized:
+                return "user"
+            return normalized.lower()
+        return "user"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        result = {
+            "message": self.message,
+            "enable_tool_calls": self.enable_tool_calls,
+            "wait_for_response": self.wait_for_response,
+            "role": self.role,
+            "request_response_format": self.request_response_format,
+            "correlationId": self.correlation_id,
+            "options": self.options,
+        }
+        if self.response_format:
+            result["response_format"] = serialize_response_format(self.response_format)
+        if self.created_at:
+            result["created_at"] = self.created_at.isoformat()
+        if self.orchestration_id:
+            result["orchestrationId"] = self.orchestration_id
+        return result
+
+    @classmethod
+    def from_json(cls, data: str) -> RunRequest:
+        """Create RunRequest from JSON string."""
+        try:
+            dict_data = json.loads(data)
+        except json.JSONDecodeError as e:
+            raise ValueError("The durable agent state is not valid JSON.") from e
+
+        return cls.from_dict(dict_data)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> RunRequest:
+        """Create RunRequest from dictionary."""
+        created_at = data.get("created_at")
+        if isinstance(created_at, str):
+            try:
+                created_at = datetime.fromisoformat(created_at)
+            except ValueError:
+                created_at = None
+
+        correlation_id = data.get("correlationId")
+        if not correlation_id:
+            raise ValueError("correlationId is required in RunRequest data")
+
+        options = data.get("options")
+
+        return cls(
+            message=data.get("message", ""),
+            correlation_id=correlation_id,
+            request_response_format=data.get("request_response_format", REQUEST_RESPONSE_FORMAT_TEXT),
+            role=cls.coerce_role(data.get("role")),
+            response_format=_deserialize_response_format(data.get("response_format")),
+            wait_for_response=data.get("wait_for_response", True),
+            enable_tool_calls=data.get("enable_tool_calls", True),
+            created_at=created_at,
+            orchestration_id=data.get("orchestrationId"),
+            options=cast(dict[str, Any], options) if isinstance(options, dict) else {},
+        )
+
+
+@dataclass
+class AgentSessionId:
+    """Represents an agent session identifier (name + key)."""
+
+    name: str
+    key: str
+
+    ENTITY_NAME_PREFIX: str = "dafx-"
+
+    @staticmethod
+    def to_entity_name(name: str) -> str:
+        return f"{AgentSessionId.ENTITY_NAME_PREFIX}{name}"
+
+    @staticmethod
+    def with_random_key(name: str) -> AgentSessionId:
+        return AgentSessionId(name=name, key=uuid.uuid4().hex)
+
+    @property
+    def entity_name(self) -> str:
+        return self.to_entity_name(self.name)
+
+    def __str__(self) -> str:
+        return f"@{self.name}@{self.key}"
+
+    def __repr__(self) -> str:
+        return f"AgentSessionId(name='{self.name}', key='{self.key}')"
+
+    @staticmethod
+    def parse(session_id_string: str, agent_name: str | None = None) -> AgentSessionId:
+        """Parses a string representation of an agent session ID.
+
+        Args:
+            session_id_string: A string in the form @name@key, or a plain key string
+                when agent_name is provided.
+            agent_name: Optional agent name to use instead of parsing from the string.
+                If provided, only the key portion is extracted from session_id_string
+                (for @name@key format) or the entire string is used as the key
+                (for plain strings).
+
+        Returns:
+            AgentSessionId instance
+
+        Raises:
+            ValueError: If the string format is invalid and agent_name is not provided
+        """
+        # Check if string is in @name@key format
+        if session_id_string.startswith("@") and "@" in session_id_string[1:]:
+            parts = session_id_string[1:].split("@", 1)
+            name = agent_name if agent_name is not None else parts[0]
+            return AgentSessionId(name=name, key=parts[1])
+
+        # Plain string format - only valid when agent_name is provided
+        if agent_name is not None:
+            return AgentSessionId(name=agent_name, key=session_id_string)
+
+        raise ValueError(f"Invalid agent session ID format: {session_id_string}")
+
+
+class DurableAgentSession(AgentSession):
+    """Durable agent session that tracks the owning :class:`AgentSessionId`."""
+
+    _SERIALIZED_SESSION_ID_KEY = "durable_session_id"
+
+    def __init__(
+        self,
+        *,
+        durable_session_id: AgentSessionId | None = None,
+        session_id: str | None = None,
+        service_session_id: str | None = None,
+    ) -> None:
+        super().__init__(session_id=session_id, service_session_id=service_session_id)
+        self.durable_session_id: AgentSessionId | None = durable_session_id
+
+    def to_dict(self) -> dict[str, Any]:
+        state = super().to_dict()
+        if self.durable_session_id is not None:
+            state[self._SERIALIZED_SESSION_ID_KEY] = str(self.durable_session_id)
+        return state
+
+    @classmethod
+    def from_session_id(
+        cls,
+        durable_session_id: AgentSessionId,
+        *,
+        session_id: str | None = None,
+        service_session_id: str | None = None,
+    ) -> DurableAgentSession:
+        """Create a DurableAgentSession from an AgentSessionId."""
+        return cls(
+            durable_session_id=durable_session_id,
+            session_id=session_id,
+            service_session_id=service_session_id,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> DurableAgentSession:
+        """Create a DurableAgentSession from a state dict."""
+        data = dict(data)  # defensive copy — avoid mutating caller's dict
+        session_id_value = data.pop(cls._SERIALIZED_SESSION_ID_KEY, None)
+        session = super().from_dict(data)
+        service_session_id = session.service_session_id
+        if service_session_id is not None and not isinstance(service_session_id, str):
+            raise ValueError("durable sessions require service_session_id to be a string when present")
+        durable_session_id: AgentSessionId | None = None
+        # We need to create a DurableAgentSession from the base AgentSession
+        if session_id_value is not None:
+            if not isinstance(session_id_value, str):
+                raise ValueError("durable_session_id must be a string when present in serialized state")
+            durable_session_id = AgentSessionId.parse(session_id_value)
+
+        durable_session = cls(
+            durable_session_id=durable_session_id,
+            session_id=session.session_id,
+            service_session_id=service_session_id,
+        )
+        durable_session.state.update(session.state)
+        return durable_session

python/packages/durabletask/agent_framework_durabletask/_orchestration_context.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Orchestration context wrapper for Durable Task Agent Framework.
+
+This module provides the DurableAIAgentOrchestrationContext class for use inside
+orchestration functions to interact with durable agents.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from durabletask.task import OrchestrationContext
+
+from ._executors import DurableAgentTask, OrchestrationAgentExecutor
+from ._shim import DurableAgentProvider, DurableAIAgent
+
+logger = logging.getLogger("agent_framework.durabletask")
+
+
+class DurableAIAgentOrchestrationContext(DurableAgentProvider[DurableAgentTask]):
+    """Orchestration context wrapper for interacting with durable agents internally.
+
+    This class wraps a durabletask OrchestrationContext and provides a convenient
+    interface for retrieving and executing durable agents from within orchestration
+    functions.
+
+    Example:
+        ```python
+        from durabletask import Orchestration
+        from agent_framework.azure import DurableAIAgentOrchestrationContext
+
+
+        def my_orchestration(context: OrchestrationContext):
+            # Wrap the context
+            agent_context = DurableAIAgentOrchestrationContext(context)
+
+            # Get an agent reference
+            agent = agent_context.get_agent("assistant")
+
+            # Run the agent (returns a Task to be yielded)
+            result = yield agent.run("Hello, how are you?")
+
+            return result.text
+        ```
+    """
+
+    def __init__(self, context: OrchestrationContext):
+        """Initialize the orchestration context wrapper.
+
+        Args:
+            context: The durabletask orchestration context to wrap
+        """
+        self._context = context
+        self._executor = OrchestrationAgentExecutor(self._context)
+        logger.debug("[DurableAIAgentOrchestrationContext] Initialized")
+
+    def get_agent(self, agent_name: str) -> DurableAIAgent[DurableAgentTask]:
+        """Retrieve a DurableAIAgent shim for the specified agent.
+
+        This method returns a proxy object that can be used to execute the agent
+        within an orchestration. The agent's run() method will return a Task that
+        must be yielded.
+
+        Args:
+            agent_name: Name of the agent to retrieve (without the dafx- prefix)
+
+        Returns:
+            DurableAIAgent instance that can be used to run the agent
+
+        Note:
+            Validation is deferred to execution time. The entity must be registered
+            on a worker with the name f"dafx-{agent_name}".
+        """
+        logger.debug("[DurableAIAgentOrchestrationContext] Creating agent proxy for: %s", agent_name)
+        return DurableAIAgent(self._executor, agent_name)

python/packages/durabletask/agent_framework_durabletask/_response_utils.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Shared utilities for handling AgentResponse parsing and validation."""
+
+import logging
+from typing import Any
+
+from agent_framework import AgentResponse
+from pydantic import BaseModel
+
+logger = logging.getLogger("agent_framework.durabletask")
+
+
+def load_agent_response(agent_response: AgentResponse | dict[str, Any] | None) -> AgentResponse:
+    """Convert raw payloads into AgentResponse instance.
+
+    Args:
+        agent_response: The response to convert, can be an AgentResponse, dict, or None
+
+    Returns:
+        AgentResponse: The converted response object
+
+    Raises:
+        ValueError: If agent_response is None
+        TypeError: If agent_response is an unsupported type
+    """
+    if agent_response is None:
+        raise ValueError("agent_response cannot be None")
+
+    logger.debug("[load_agent_response] Loading agent response of type: %s", type(agent_response))
+
+    if isinstance(agent_response, AgentResponse):
+        return agent_response
+    if isinstance(agent_response, dict):
+        logger.debug("[load_agent_response] Converting dict payload using AgentResponse.from_dict")
+        return AgentResponse.from_dict(agent_response)
+
+    raise TypeError(f"Unsupported type for agent_response: {type(agent_response)}")
+
+
+def ensure_response_format(
+    response_format: type[BaseModel] | None,
+    correlation_id: str,
+    response: AgentResponse,
+) -> None:
+    """Ensure the AgentResponse value is parsed into the expected response_format.
+
+    This function modifies the response in-place by parsing its value attribute
+    into the specified Pydantic model format.
+
+    Args:
+        response_format: Optional Pydantic model class to parse the response value into
+        correlation_id: Correlation ID for logging purposes
+        response: The AgentResponse object to validate and parse
+
+    Raises:
+        ValueError: If response_format is specified but response.value cannot be parsed
+    """
+    if response_format is not None:
+        # Set the response format on the response so .value knows how to parse
+        response._response_format = response_format  # pyright: ignore[reportPrivateUsage]
+        response._value_parsed = False  # pyright: ignore[reportPrivateUsage]  # Reset to allow re-parsing with new format
+
+        # Access response.value to trigger parsing (may raise ValidationError)
+        # Validate that parsing succeeded
+        if not isinstance(response.value, response_format):
+            raise ValueError(
+                f"Response value could not be parsed into required format {response_format.__name__} "
+                f"for correlation_id {correlation_id}"
+            )
+
+        logger.debug(
+            "[ensure_response_format] Loaded AgentResponse.value for correlation_id %s with type: %s",
+            correlation_id,
+            type(response.value).__name__,
+        )

python/packages/durabletask/agent_framework_durabletask/_shim.py

@@ -0,0 +1,168 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Durable Agent Shim for Durable Task Framework.
+
+This module provides the DurableAIAgent shim that implements SupportsAgentRun
+and provides a consistent interface for both Client and Orchestration contexts.
+The actual execution is delegated to the context-specific providers.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Generic, Literal, TypeVar
+
+from agent_framework import AgentSession, ServiceSessionId, SupportsAgentRun, normalize_messages
+from agent_framework._types import AgentRunInputs
+
+from ._executors import DurableAgentExecutor
+from ._models import DurableAgentSession
+
+# TypeVar for the task type returned by executors
+# Covariant because TaskT only appears in return positions (output)
+TaskT = TypeVar("TaskT", covariant=True)
+
+
+class DurableAgentProvider(ABC, Generic[TaskT]):
+    """Abstract provider for constructing durable agent proxies.
+
+    Implemented by context-specific wrappers (client/orchestration) to return a
+    `DurableAIAgent` shim backed by their respective `DurableAgentExecutor`
+    implementation, ensuring a consistent `get_agent` entry point regardless of
+    execution context.
+    """
+
+    @abstractmethod
+    def get_agent(self, agent_name: str) -> DurableAIAgent[TaskT]:
+        """Retrieve a DurableAIAgent shim for the specified agent.
+
+        Args:
+            agent_name: Name of the agent to retrieve
+
+        Returns:
+            DurableAIAgent instance that can be used to run the agent
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses
+        """
+        raise NotImplementedError("Subclasses must implement get_agent()")
+
+
+class DurableAIAgent(SupportsAgentRun, Generic[TaskT]):
+    """A durable agent proxy that delegates execution to the provider.
+
+    This class implements SupportsAgentRun but with one critical difference:
+    - SupportsAgentRun.run() returns a Coroutine (async, must await)
+    - DurableAIAgent.run() returns TaskT (sync Task object - must yield
+        or the AgentResponse directly in the case of TaskHubGrpcClient)
+
+    This represents fundamentally different execution models but maintains the same
+    interface contract for all other properties and methods.
+
+    The underlying provider determines how execution occurs (entity calls, HTTP requests, etc.)
+    and what type of Task object is returned.
+
+    Type Parameters:
+        TaskT: The task type returned by this agent (e.g., AgentResponse, DurableAgentTask, AgentTask)
+    """
+
+    id: str
+    name: str
+    display_name: str
+    description: str | None
+
+    def __init__(self, executor: DurableAgentExecutor[TaskT], name: str, *, agent_id: str | None = None):
+        """Initialize the shim with a provider and agent name.
+
+        Args:
+            executor: The execution provider (Client or OrchestrationContext)
+            name: The name of the agent to execute
+            agent_id: Optional unique identifier for the agent (defaults to name)
+        """
+        self._executor = executor
+        self.name = name  # pyright: ignore[reportIncompatibleVariableOverride]
+        self.id = agent_id if agent_id is not None else name
+        self.display_name = name
+        self.description = f"Durable agent proxy for {name}"
+
+    def run(  # type: ignore[override]
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: Literal[False] = False,
+        session: AgentSession | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> TaskT:
+        """Execute the agent via the injected provider.
+
+        Args:
+            messages: The message(s) to send to the agent
+            stream: Whether to use streaming for the response (must be False)
+                DurableAgents do not support streaming mode.
+            session: Optional agent session for conversation context
+            options: Optional options dictionary. Supported keys include
+                ``response_format``, ``enable_tool_calls``, and ``wait_for_response``.
+                Additional keys are forwarded to the agent execution.
+
+        Note:
+            This method overrides SupportsAgentRun.run() with a different return type:
+            - SupportsAgentRun.run() returns Coroutine[Any, Any, AgentResponse] (async)
+            - DurableAIAgent.run() returns TaskT (Task object for yielding)
+
+            This is intentional to support orchestration contexts that use yield patterns
+            instead of async/await patterns.
+
+        Returns:
+            TaskT: The task type specific to the executor
+
+        Raises:
+            ValueError: If wait_for_response=False is used in an unsupported context
+        """
+        if stream is not False:
+            raise ValueError("DurableAIAgent does not support streaming mode (stream must be False)")
+        message_str = self._normalize_messages(messages)
+
+        run_request = self._executor.get_run_request(
+            message=message_str,
+            options=options,
+        )
+
+        return self._executor.run_durable_agent(
+            agent_name=self.name,
+            run_request=run_request,
+            session=session,
+        )
+
+    def create_session(self, *, session_id: str | None = None) -> DurableAgentSession:
+        """Create a new agent session via the provider."""
+        return self._executor.get_new_session(self.name)
+
+    def get_session(self, service_session_id: str | ServiceSessionId, *, session_id: str | None = None) -> AgentSession:
+        """Retrieve an existing session via the provider."""
+        if not isinstance(service_session_id, str):
+            raise ValueError("DurableAIAgent requires service_session_id to be a string")
+        return self._executor.get_new_session(self.name, service_session_id=service_session_id, session_id=session_id)
+
+    def _normalize_messages(self, messages: AgentRunInputs | None) -> str:
+        """Convert supported message inputs to a single string.
+
+        Args:
+            messages: The messages to normalize
+
+        Returns:
+            A single string representation of the messages
+
+        Raises:
+            ValueError: If normalized messages contain non-text content only.
+        """
+        normalized_messages = normalize_messages(messages)
+        if not normalized_messages:
+            return ""
+
+        message_texts: list[str] = []
+        for message in normalized_messages:
+            if not message.text:
+                raise ValueError("DurableAIAgent only supports text message inputs.")
+            message_texts.append(message.text)
+
+        return "\n".join(message_texts)

python/packages/durabletask/agent_framework_durabletask/_worker.py

@@ -0,0 +1,421 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Worker wrapper for Durable Task Agent Framework.
+
+This module provides the DurableAIAgentWorker class that wraps a durabletask worker
+and enables registration of agents as durable entities, and optionally workflows
+as durable orchestrations with automatically generated activity functions.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from agent_framework import SupportsAgentRun, Workflow
+from durabletask.task import ActivityContext, OrchestrationContext
+from durabletask.worker import TaskHubGrpcWorker
+
+from ._async_bridge import run_agent_coroutine
+from ._callbacks import AgentResponseCallbackProtocol
+from ._entities import AgentEntity, DurableTaskEntityStateProvider
+from ._workflows.activity import execute_workflow_activity
+from ._workflows.dt_context import DurableTaskWorkflowContext
+from ._workflows.naming import (
+    validate_executor_id,
+    validate_workflow_name,
+    workflow_executor_activity_name,
+    workflow_orchestrator_name,
+    workflow_scoped_executor_id,
+)
+from ._workflows.orchestrator import run_workflow_orchestrator
+from ._workflows.registration import collect_hosted_workflows, plan_workflow_registration
+
+logger = logging.getLogger("agent_framework.durabletask")
+
+
+class DurableAIAgentWorker:
+    """Wrapper for a durabletask worker that hosts agents and workflows.
+
+    This class wraps an existing TaskHubGrpcWorker instance and is the single
+    host-side registration surface for a worker process. It supports two
+    complementary kinds of work:
+
+    - **Agents** via :meth:`add_agent`, which registers each agent as a durable entity.
+    - **Workflows** via :meth:`configure_workflow`, which registers a MAF
+      ``Workflow`` (its agent executors as entities, its non-agent executors as
+      activities, and the workflow orchestrator).
+
+    A single worker process commonly hosts both, so registration is intentionally
+    aggregated on one object rather than split per kind. (On the *client* side the
+    surfaces are split into :class:`DurableAIAgentClient` and ``DurableWorkflowClient``,
+    because a caller invokes one or the other.)
+
+    Example:
+        ```python
+        from durabletask.worker import TaskHubGrpcWorker
+        from agent_framework import Agent
+        from agent_framework.openai import OpenAIChatCompletionClient
+        from agent_framework_durabletask import DurableAIAgentWorker
+
+        # Create the underlying worker
+        worker = TaskHubGrpcWorker(host_address="localhost:4001")
+
+        # Wrap it with the agent worker
+        agent_worker = DurableAIAgentWorker(worker)
+
+        # Register agents (or call configure_workflow(workflow) to host a workflow)
+        client = OpenAIChatCompletionClient()
+        my_agent = Agent(client=client, name="assistant")
+        agent_worker.add_agent(my_agent)
+
+        # Start the worker
+        worker.start()
+        ```
+    """
+
+    def __init__(
+        self,
+        worker: TaskHubGrpcWorker,
+        callback: AgentResponseCallbackProtocol | None = None,
+    ):
+        """Initialize the worker wrapper.
+
+        Args:
+            worker: The durabletask worker instance to wrap
+            callback: Optional callback for agent response notifications
+        """
+        self._worker = worker
+        self._callback = callback
+        self._registered_agents: dict[str, SupportsAgentRun] = {}
+        self._workflows: dict[str, Workflow] = {}
+        # Every workflow whose orchestration has been registered (top-level plus nested
+        # sub-workflows), keyed by case-folded name -> the registered instance, so a
+        # sub-workflow shared across the tree is registered once while two different
+        # workflows whose names collide (including case-only differences) are rejected.
+        self._registered_orchestrations: dict[str, Workflow] = {}
+        logger.debug("[DurableAIAgentWorker] Initialized with worker type: %s", type(worker).__name__)
+
+    def add_agent(
+        self,
+        agent: SupportsAgentRun,
+        callback: AgentResponseCallbackProtocol | None = None,
+        *,
+        entity_id: str | None = None,
+    ) -> None:
+        """Register an agent with the worker.
+
+        This method creates a durable entity class for the agent and registers
+        it with the underlying durabletask worker. The entity will be accessible
+        by the name "dafx-{entity_id or agent_name}".
+
+        Args:
+            agent: The agent to register (must have a name)
+            callback: Optional callback for this specific agent (overrides worker-level callback)
+            entity_id: Optional identity to register the entity under instead of
+                ``agent.name``. Workflow hosting passes the executor's ``id`` so the
+                entity matches the identity the orchestrator dispatches to.
+
+        Raises:
+            ValueError: If the agent doesn't have a name or is already registered
+        """
+        registration_name = entity_id or agent.name
+        if not registration_name:
+            raise ValueError("Agent must have a name to be registered")
+
+        if registration_name in self._registered_agents:
+            raise ValueError(f"Agent '{registration_name}' is already registered")
+
+        logger.info(
+            "[DurableAIAgentWorker] Registering agent: %s as entity: dafx-%s", registration_name, registration_name
+        )
+
+        # Store the agent reference
+        self._registered_agents[registration_name] = agent
+
+        # Use agent-specific callback if provided, otherwise use worker-level callback
+        effective_callback = callback or self._callback
+
+        # Create a configured entity class using the factory
+        entity_class = self.__create_agent_entity(agent, effective_callback, entity_id=registration_name)
+
+        # Register the entity class with the worker
+        # The worker.add_entity method takes a class
+        entity_registered: str = self._worker.add_entity(entity_class)
+
+        logger.debug(
+            "[DurableAIAgentWorker] Successfully registered entity class %s for agent: %s",
+            entity_registered,
+            registration_name,
+        )
+
+    def start(self) -> None:
+        """Start the worker to begin processing tasks.
+
+        Note:
+            This method delegates to the underlying worker's start method.
+            The worker will block until stopped.
+        """
+        logger.info("[DurableAIAgentWorker] Starting worker with %d registered agents", len(self._registered_agents))
+        self._worker.start()
+
+    def stop(self) -> None:
+        """Stop the worker gracefully.
+
+        Note:
+            This method delegates to the underlying worker's stop method.
+        """
+        logger.info("[DurableAIAgentWorker] Stopping worker")
+        self._worker.stop()
+
+    @property
+    def registered_agent_names(self) -> list[str]:
+        """Get the names of all registered agents.
+
+        Returns:
+            List of agent names (without the dafx- prefix)
+        """
+        return list(self._registered_agents.keys())
+
+    @property
+    def registered_workflow_names(self) -> list[str]:
+        """Get the names of all workflows configured on this worker.
+
+        Returns:
+            List of workflow names (the identities used to derive each workflow's
+            ``dafx-{name}`` orchestration).
+        """
+        return list(self._workflows.keys())
+
+    # -----------------------------------------------------------------
+    # Workflow support
+    # -----------------------------------------------------------------
+
+    def configure_workflow(
+        self,
+        workflow: Workflow,
+        callback: AgentResponseCallbackProtocol | None = None,
+    ) -> None:
+        """Register a :class:`Workflow` for automatic orchestration.
+
+        This extracts agents from the workflow and registers them as durable
+        entities, registers non-agent executors as activities, and creates an
+        orchestrator function that drives the workflow graph.
+
+        Multiple workflows can be hosted on one worker: call this method once per
+        workflow. Each workflow is keyed by its :attr:`Workflow.name`, and its
+        durable primitives are scoped by that name (orchestration
+        ``dafx-{name}``; activities/entities ``dafx-{name}-{executorId}``) so two
+        co-hosted workflows that reuse an executor id do not collide.
+
+        Sub-workflows nest: if the workflow contains
+        :class:`~agent_framework.WorkflowExecutor` nodes, each inner workflow's
+        orchestration/agents/activities are registered too (deduped by name) so the
+        parent can drive them as durable child orchestrations.
+
+        Args:
+            workflow: The MAF :class:`Workflow` to register. Must have an explicit,
+                stable :attr:`Workflow.name` (an auto-generated
+                ``WorkflowBuilder-<uuid>`` name is rejected because it is not stable
+                across restarts and would break durable resume). Every nested
+                sub-workflow must likewise be named.
+            callback: Optional callback for agent response notifications.
+
+        Raises:
+            ValueError: If the workflow (or a nested sub-workflow) name is missing,
+                invalid, or auto-generated, or if the top-level workflow name is
+                already registered on this worker.
+        """
+        workflow_name = workflow.name
+        validate_workflow_name(workflow_name)
+        if any(name.casefold() == workflow_name.casefold() for name in self._workflows):
+            raise ValueError(
+                f"Workflow '{workflow_name}' is already registered on this worker "
+                "(workflow names are compared case-insensitively)."
+            )
+
+        # Validate the whole composition (top-level plus every nested sub-workflow)
+        # up front, so an invalid/auto-generated nested name (or an executor id that
+        # would break durable naming / nested-HITL addressing) fails before any
+        # registration side effects leave the worker partially configured.
+        hosted_workflows = list(collect_hosted_workflows(workflow))
+        for hosted in hosted_workflows:
+            validate_workflow_name(hosted.name)
+            for executor_id in hosted.executors:
+                validate_executor_id(executor_id)
+
+        # Check every cross-call collision *before* mutating any state, so a clash
+        # between a nested sub-workflow and an already-registered orchestration cannot
+        # leave the worker partially configured (e.g. the top-level name added to
+        # ``_workflows`` while a later child fails). Registration below is then a pure
+        # commit step.
+        for hosted in hosted_workflows:
+            existing = self._registered_orchestrations.get(hosted.name.casefold())
+            if existing is not None and existing is not hosted:
+                raise ValueError(
+                    f"A different workflow named '{hosted.name}' collides with already-registered "
+                    f"'{existing.name}' on this worker. A workflow name maps to a single durable "
+                    f"orchestration ('dafx-{hosted.name}'), compared case-insensitively; rename one "
+                    "of them."
+                )
+
+        self._workflows[workflow_name] = workflow
+
+        # Commit: register the top-level workflow and every nested sub-workflow (deduped
+        # by name), so the parent can drive sub-workflows as durable child orchestrations.
+        for hosted in hosted_workflows:
+            if hosted.name.casefold() in self._registered_orchestrations:
+                continue
+            self._register_single_workflow(hosted, callback)
+
+    def _register_single_workflow(
+        self,
+        workflow: Workflow,
+        callback: AgentResponseCallbackProtocol | None,
+    ) -> None:
+        """Register one workflow's durable primitives (no recursion into sub-workflows).
+
+        The "what to register" decision (agent -> entity, non-agent -> activity,
+        sub-workflow -> child orchestration) is shared with the Azure Functions host
+        via ``plan_workflow_registration``.
+        """
+        validate_workflow_name(workflow.name)
+        self._registered_orchestrations[workflow.name.casefold()] = workflow
+        plan = plan_workflow_registration(workflow)
+
+        # Register agent executors as durable entities, scoped by workflow name so
+        # two workflows that reuse an executor id register distinct entities. The
+        # entity is keyed by the scoped identity (the same identity the orchestrator
+        # dispatches to); the entity *key* at run time is the orchestration instance
+        # id, which keeps conversation state isolated per run.
+        for agent_executor in plan.agent_executors:
+            scoped_id = workflow_scoped_executor_id(workflow.name, agent_executor.id)
+            if scoped_id not in self._registered_agents:
+                self.add_agent(agent_executor.agent, callback=callback, entity_id=scoped_id)
+
+        # Register non-agent executors as durable activities, scoped by workflow name.
+        # WorkflowExecutor nodes are intentionally not registered as activities: their
+        # inner workflows are registered separately (above, via collect_hosted_workflows)
+        # and driven as child orchestrations.
+        for executor in plan.activity_executors:
+            self._register_executor_activity(workflow, executor)
+
+        # Register this workflow's orchestrator under its per-workflow name.
+        self._register_workflow_orchestrator(workflow)
+
+        logger.info(
+            "[DurableAIAgentWorker] Workflow '%s' configured with %d executors "
+            "(%d agents, %d activities, %d sub-workflows)",
+            workflow.name,
+            len(workflow.executors),
+            len(plan.agent_executors),
+            len(plan.activity_executors),
+            len(plan.subworkflow_executors),
+        )
+
+    def _register_executor_activity(self, workflow: Workflow, executor: Any) -> None:
+        """Register a non-agent executor as a durabletask activity (workflow-scoped)."""
+        captured_executor = executor
+        captured_workflow = workflow
+        activity_name = workflow_executor_activity_name(workflow.name, executor.id)
+
+        def executor_activity(ctx: ActivityContext, input_data: str) -> str:
+            return execute_workflow_activity(captured_executor, input_data, captured_workflow)
+
+        # Give the function the expected name for registration
+        executor_activity.__name__ = activity_name
+        executor_activity.__qualname__ = activity_name
+
+        self._worker.add_activity(executor_activity)
+        logger.debug("[DurableAIAgentWorker] Registered activity: %s", activity_name)
+
+    def _register_workflow_orchestrator(self, workflow: Workflow) -> None:
+        """Register a workflow's orchestrator function under its per-workflow name."""
+        captured_workflow = workflow
+        orchestrator_name = workflow_orchestrator_name(workflow.name)
+
+        def workflow_orchestrator(context: OrchestrationContext, input_data: Any) -> Any:
+            # Pass the deserialized client input straight to the shared engine, which
+            # reconstructs the start executor's declared type (see _coerce_initial_input).
+            initial_message = input_data
+            shared_state: dict[str, Any] = {}
+
+            dt_ctx = DurableTaskWorkflowContext(context)
+            outputs = yield from run_workflow_orchestrator(dt_ctx, captured_workflow, initial_message, shared_state)
+            return outputs  # noqa: B901
+
+        workflow_orchestrator.__name__ = orchestrator_name
+        workflow_orchestrator.__qualname__ = orchestrator_name
+
+        self._worker.add_orchestrator(workflow_orchestrator)
+        logger.debug("[DurableAIAgentWorker] Registered workflow orchestrator: %s", orchestrator_name)
+
+    def __create_agent_entity(
+        self,
+        agent: SupportsAgentRun,
+        callback: AgentResponseCallbackProtocol | None = None,
+        *,
+        entity_id: str | None = None,
+    ) -> type[DurableTaskEntityStateProvider]:
+        """Factory function to create a DurableEntity class configured with an agent.
+
+        This factory creates a new class that combines the entity state provider
+        with the agent execution logic. Each agent gets its own entity class.
+
+        Args:
+            agent: The agent instance to wrap
+            callback: Optional callback for agent responses
+            entity_id: Optional identity to register the entity under instead of
+                ``agent.name`` (used by workflow hosting to key entities by
+                executor id).
+
+        Returns:
+            A new DurableEntity subclass configured for this agent
+        """
+        agent_name = entity_id or agent.name or type(agent).__name__
+        entity_name = f"dafx-{agent_name}"
+
+        class ConfiguredAgentEntity(DurableTaskEntityStateProvider):
+            """Durable entity configured with a specific agent instance."""
+
+            def __init__(self) -> None:
+                super().__init__()
+                # Create the AgentEntity with this state provider
+                self._agent_entity = AgentEntity(
+                    agent=agent,
+                    callback=callback,
+                    state_provider=self,
+                )
+                logger.debug(
+                    "[ConfiguredAgentEntity] Initialized entity for agent: %s (entity name: %s)",
+                    agent_name,
+                    entity_name,
+                )
+
+            def run(self, request: Any) -> Any:
+                """Handle run requests from clients or orchestrations.
+
+                Args:
+                    request: RunRequest as dict or string
+
+                Returns:
+                    AgentResponse as dict
+                """
+                logger.debug("[ConfiguredAgentEntity.run] Executing agent: %s", agent_name)
+                # Run on the shared persistent loop so async resources created by
+                # shared agent clients/credentials stay bound to a live loop across
+                # successive entity invocations (avoids cross-loop hangs).
+                response = run_agent_coroutine(self._agent_entity.run(request))
+                return response.to_dict()
+
+            def reset(self) -> None:
+                """Reset the agent's conversation history."""
+                logger.debug("[ConfiguredAgentEntity.reset] Resetting agent: %s", agent_name)
+                self._agent_entity.reset()
+
+        # Set the entity name to match the prefixed agent name
+        # This is used by durabletask to register the entity
+        ConfiguredAgentEntity.__name__ = entity_name
+        ConfiguredAgentEntity.__qualname__ = entity_name
+
+        return ConfiguredAgentEntity

python/packages/durabletask/agent_framework_durabletask/_workflows/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Durable hosting of Microsoft Agent Framework workflows.
+
+This subpackage turns a MAF :class:`~agent_framework.Workflow` into durable
+primitives -- a single orchestrator, agent entities, and non-agent executor
+activities -- that run on either a standalone Durable Task worker or Azure
+Functions. The host-agnostic engine lives here; each host programs against the
+:class:`~.context.WorkflowOrchestrationContext` protocol.
+"""

python/packages/durabletask/agent_framework_durabletask/_workflows/activity.py

@@ -0,0 +1,186 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Host-agnostic execution of non-agent workflow executors as durable activities.
+
+When a MAF :class:`Workflow` runs as a durable orchestration, each non-agent
+executor is dispatched as a durable *activity*. The activity body is identical
+regardless of host (Azure Functions or a standalone durabletask worker): it
+deserializes the activity input, runs the executor (or a human-in-the-loop
+response handler), diffs the shared state, and serializes the executor's
+outputs, sent messages, shared-state changes, and any pending HITL requests back
+to the orchestrator.
+
+This module provides that shared body as :func:`execute_workflow_activity` so
+both host adapters call one implementation instead of duplicating it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from copy import deepcopy
+from typing import Any, cast
+
+from agent_framework import Executor, Workflow, WorkflowEvent
+from agent_framework._workflows._runner_context import YieldOutputEventType
+from agent_framework._workflows._state import State
+
+from .orchestrator import (
+    SOURCE_HITL_RESPONSE,
+    SOURCE_ORCHESTRATOR,
+    execute_hitl_response_handler,
+)
+from .runner_context import CapturingRunnerContext
+from .serialization import deserialize_value, serialize_value, serialize_workflow_event
+
+
+def execute_workflow_activity(executor: Executor, input_json: str, workflow: Workflow | None = None) -> str:
+    """Execute a single non-agent workflow executor and return its serialized result.
+
+    This is the host-agnostic activity body shared by the Azure Functions and
+    standalone durabletask workflow hosts.
+
+    Args:
+        executor: The non-agent executor instance to run.
+        input_json: JSON-encoded activity input with keys ``message``,
+            ``shared_state_snapshot``, and ``source_executor_ids``.
+        workflow: The owning workflow, used to classify the executor's
+            ``yield_output`` payloads as final ``output`` vs ``intermediate``.
+            When omitted, all yielded outputs are treated as final outputs.
+
+    Returns:
+        A JSON string with keys ``sent_messages``, ``outputs``, ``events``,
+        ``shared_state_updates``, ``shared_state_deletes``, and
+        ``pending_request_info_events``.
+
+    Raises:
+        ValueError: If the input does not decode to a JSON object, or a HITL
+            message payload is not a JSON object.
+    """
+    data_obj = json.loads(input_json)
+    if not isinstance(data_obj, dict):
+        raise ValueError("Activity input must decode to a JSON object")
+    data = cast(dict[str, Any], data_obj)
+
+    message_data = data.get("message")
+    # The orchestrator may pass null for these when shared state / sources are
+    # omitted, so coerce None to the appropriate empty default.
+    shared_state_snapshot: dict[str, Any] = data.get("shared_state_snapshot") or {}
+    source_executor_ids = cast(list[str], data.get("source_executor_ids") or [SOURCE_ORCHESTRATOR])
+
+    # Reconstruct the message - deserialize_value restores the original typed
+    # objects from the encoded data (with type markers).
+    message = deserialize_value(message_data)
+
+    # A HITL response is identified by a source id starting with the HITL prefix.
+    is_hitl_response = any(s.startswith(SOURCE_HITL_RESPONSE) for s in source_executor_ids)
+
+    def classify_yielded_output(executor_id: str) -> YieldOutputEventType | None:
+        # Mirror the core runner's classification so intermediate executors'
+        # yields are not surfaced as final workflow outputs.
+        if workflow is None:
+            return "output"
+        if workflow.is_terminal_executor(executor_id):
+            return "output"
+        if workflow.is_intermediate_executor(executor_id):
+            return "intermediate"
+        return None
+
+    async def _run() -> dict[str, Any]:
+        runner_context = CapturingRunnerContext()
+        runner_context.set_yield_output_classifier(classify_yielded_output)
+        shared_state = State()
+
+        # Deserialize shared state values to reconstruct dataclasses / Pydantic models.
+        deserialized_state: dict[str, Any] = {str(k): deserialize_value(v) for k, v in shared_state_snapshot.items()}
+        # Snapshot the deserialized (in-memory) state for diffing. State.export_state()
+        # returns the in-memory committed objects, so the snapshot must hold objects
+        # too (deepcopy) - comparing against a serialized snapshot would mark every
+        # key as changed.
+        original_snapshot = deepcopy(deserialized_state)
+        shared_state.import_state(deserialized_state)
+
+        if is_hitl_response:
+            if not isinstance(message_data, dict):
+                raise ValueError("HITL message payload must be a JSON object")
+            await execute_hitl_response_handler(
+                executor=executor,
+                hitl_message=cast(dict[str, Any], message_data),
+                shared_state=shared_state,
+                runner_context=runner_context,
+            )
+        else:
+            await executor.execute(
+                message=message,
+                source_executor_ids=source_executor_ids,
+                state=shared_state,
+                runner_context=runner_context,
+            )
+
+        # Commit pending state changes and compute the diff vs the original snapshot.
+        shared_state.commit()
+        current_state = shared_state.export_state()
+        original_keys: set[str] = set(original_snapshot.keys())
+        current_keys: set[str] = set(current_state.keys())
+
+        # Deleted = was in original, not in current.
+        deletes: set[str] = original_keys - current_keys
+
+        # Updates = keys that are new or whose value changed.
+        updates: dict[str, Any] = {}
+        for key in current_keys:
+            if key not in original_keys or current_state[key] != original_snapshot.get(key):
+                updates[key] = current_state[key]
+
+        sent_messages = await runner_context.drain_messages()
+        events = await runner_context.drain_events()
+
+        # Serialize the executor's workflow events so the orchestrator can republish
+        # them to the streaming custom status. Output payloads are also extracted
+        # separately for message routing and the final workflow result.
+        outputs: list[Any] = []
+        serialized_events: list[dict[str, Any]] = []
+        for event in events:
+            if not isinstance(event, WorkflowEvent):
+                continue
+            serialized_events.append(serialize_workflow_event(event))
+            if event.type == "output":
+                outputs.append(serialize_value(event.data))
+
+        # Serialize pending HITL request info events for the orchestrator.
+        pending_request_info_events = await runner_context.get_pending_request_info_events()
+        serialized_pending_requests: list[dict[str, Any]] = []
+        for _request_id, event in pending_request_info_events.items():
+            serialized_pending_requests.append({
+                "request_id": event.request_id,
+                "source_executor_id": event.source_executor_id,
+                "data": serialize_value(event.data),
+                "request_type": f"{type(event.data).__module__}:{type(event.data).__name__}",
+                "response_type": f"{event.response_type.__module__}:{event.response_type.__name__}"
+                if event.response_type
+                else None,
+            })
+
+        # Serialize sent messages for JSON compatibility.
+        serialized_sent_messages: list[dict[str, Any]] = []
+        for _source_id, msg_list in sent_messages.items():
+            for msg in msg_list:
+                serialized_sent_messages.append({
+                    "message": serialize_value(msg.data),
+                    "target_id": msg.target_id,
+                    "source_id": msg.source_id,
+                })
+
+        serialized_updates = {k: serialize_value(v) for k, v in updates.items()}
+
+        return {
+            "sent_messages": serialized_sent_messages,
+            "outputs": outputs,
+            "events": serialized_events,
+            "shared_state_updates": serialized_updates,
+            "shared_state_deletes": list(deletes),
+            "pending_request_info_events": serialized_pending_requests,
+        }
+
+    result = asyncio.run(_run())
+    return json.dumps(result)

python/packages/durabletask/agent_framework_durabletask/_workflows/client.py

@@ -0,0 +1,527 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Workflow client wrapper for Durable Task Agent Framework.
+
+This module provides :class:`DurableWorkflowClient` for external clients to start,
+await, and drive (including human-in-the-loop) workflows registered on a worker via
+``DurableAIAgentWorker.configure_workflow``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+from collections.abc import AsyncIterator
+from typing import Any, cast
+
+from agent_framework import WorkflowEvent
+from durabletask.client import TaskHubGrpcClient
+
+from .naming import (
+    qualify_subworkflow_request_id,
+    split_subworkflow_request_id,
+    workflow_orchestrator_name,
+)
+from .serialization import (
+    deserialize_workflow_event,
+    deserialize_workflow_output,
+    strip_pickle_markers,
+    strip_subworkflow_markers,
+)
+
+logger = logging.getLogger("agent_framework.durabletask")
+
+
+class DurableWorkflowClient:
+    """Client wrapper for starting and driving durable workflows externally.
+
+    This class wraps a durabletask ``TaskHubGrpcClient`` and provides a convenient
+    interface for the workflow registered by ``DurableAIAgentWorker.configure_workflow``:
+    starting it, awaiting its output, and responding to human-in-the-loop (HITL) pauses.
+
+    For interacting with individual durable *agents*, use
+    :class:`~agent_framework_durabletask.DurableAIAgentClient` instead. Both wrap the
+    same underlying ``TaskHubGrpcClient``, so an application that needs both can
+    construct both over one client.
+
+    Example:
+        ```python
+        from durabletask.azuremanaged.client import DurableTaskSchedulerClient
+        from agent_framework.azure import DurableWorkflowClient
+
+        # Create the underlying client
+        client = DurableTaskSchedulerClient(host_address="localhost:8080", taskhub="default")
+
+        # Wrap it with the workflow client, defaulting to the workflow named "orders"
+        workflow_client = DurableWorkflowClient(client, workflow_name="orders")
+
+        # Start a workflow and wait for its output
+        instance_id = workflow_client.start_workflow(input="some input")
+        output = workflow_client.await_workflow_output(instance_id)
+        print(output)
+
+        # A client without a default targets workflows explicitly per call:
+        multi = DurableWorkflowClient(client)
+        instance_id = multi.start_workflow(input="...", workflow_name="billing")
+        ```
+    """
+
+    def __init__(self, client: TaskHubGrpcClient, *, workflow_name: str | None = None):
+        """Initialize the workflow client wrapper.
+
+        Args:
+            client: The durabletask client instance to wrap.
+            workflow_name: Optional default workflow name to target. When set, the
+                per-call ``workflow_name`` may be omitted. When a worker hosts a
+                single workflow, set this once here; when it hosts several, either
+                set a default and override per call, or pass ``workflow_name`` on
+                each call.
+        """
+        self._client = client
+        self._default_workflow_name = workflow_name
+        logger.debug("[DurableWorkflowClient] Initialized with client type: %s", type(client).__name__)
+
+    def _resolve_workflow_name(self, workflow_name: str | None) -> str:
+        """Resolve the effective workflow name from a per-call value or the default.
+
+        Raises:
+            ValueError: If neither a per-call ``workflow_name`` nor a constructor
+                default was provided.
+        """
+        name = workflow_name or self._default_workflow_name
+        if not name:
+            raise ValueError(
+                "No workflow name provided. Pass workflow_name=... (or set a default on "
+                "DurableWorkflowClient(workflow_name=...)) so the client can target the "
+                "right orchestration."
+            )
+        return name
+
+    def start_workflow(
+        self, input: Any = None, *, workflow_name: str | None = None, instance_id: str | None = None
+    ) -> str:
+        """Start the workflow orchestration registered by ``configure_workflow``.
+
+        This schedules the orchestration ``dafx-{workflow_name}`` that
+        ``DurableAIAgentWorker.configure_workflow`` auto-registers, so callers do
+        not need to know its internal name.
+
+        Args:
+            input: The initial message/payload for the workflow.
+            workflow_name: The workflow to start. Optional if a default was set on
+                the client; required otherwise.
+            instance_id: Optional explicit orchestration instance ID. If omitted, one
+                is generated.
+
+        Returns:
+            The orchestration instance ID, for use with ``await_workflow_output``.
+        """
+        orchestration_name = workflow_orchestrator_name(self._resolve_workflow_name(workflow_name))
+        new_instance_id = self._client.schedule_new_orchestration(
+            orchestration_name,
+            # Neutralize a forged sub-workflow envelope before scheduling: only an
+            # internal child dispatch (post trust boundary) may carry those reserved
+            # keys, so stripping them here keeps untrusted input off the orchestrator's
+            # trusted-deserialization path even if start_workflow is exposed remotely.
+            input=strip_subworkflow_markers(input),
+            instance_id=instance_id,
+        )
+        logger.debug("[DurableWorkflowClient] Started workflow instance: %s", new_instance_id)
+        return new_instance_id
+
+    def _is_owned_orchestration(self, state: Any, workflow_name: str | None) -> bool:
+        """Return whether ``state`` belongs to the targeted workflow.
+
+        Ownership validation is opt-in: when neither a per-call ``workflow_name``
+        nor a constructor default is set there is nothing to validate against, so
+        this returns ``True``. When a name is resolvable, the instance's
+        orchestration name must equal ``dafx-{workflow_name}`` (compared
+        case-insensitively, mirroring the Azure Functions host's route-scoping
+        check). This guards against addressing an instance that belongs to a
+        different workflow on the same task hub.
+        """
+        name = workflow_name or self._default_workflow_name
+        if not name:
+            return True
+        expected = workflow_orchestrator_name(name)
+        actual = getattr(state, "name", None)
+        return isinstance(actual, str) and actual.casefold() == expected.casefold()
+
+    def await_workflow_output(
+        self, instance_id: str, *, workflow_name: str | None = None, timeout_seconds: int = 300
+    ) -> Any:
+        """Wait for a workflow orchestration to complete and return its output.
+
+        Args:
+            instance_id: The instance ID returned by ``start_workflow``.
+            workflow_name: Optional workflow name; when set (or a client default is
+                set) the instance's orchestration is validated to belong to that
+                workflow.
+            timeout_seconds: Maximum time, in seconds, to wait for completion.
+
+        Returns:
+            The deserialized workflow output (typically a list of yielded outputs),
+            or ``None`` if the workflow produced no output.
+
+        Raises:
+            TimeoutError: If the workflow does not complete within ``timeout_seconds``.
+            RuntimeError: If the workflow completes with a non-successful status.
+            ValueError: If the instance does not belong to the targeted workflow.
+        """
+        metadata = self._client.wait_for_orchestration_completion(instance_id, timeout=timeout_seconds)
+        if metadata is None:
+            raise TimeoutError(f"Workflow '{instance_id}' did not complete within {timeout_seconds}s")
+
+        if not self._is_owned_orchestration(metadata, workflow_name):
+            raise ValueError(f"Instance '{instance_id}' does not belong to the targeted workflow.")
+
+        status = metadata.runtime_status.name
+        if status != "COMPLETED":
+            raise RuntimeError(f"Workflow '{instance_id}' ended with status {status}: {metadata.serialized_output}")
+
+        if metadata.serialized_output is None:
+            return None
+        # The shared activity encodes each yielded output with serialize_value()
+        # before it reaches the orchestrator, so typed objects come back as
+        # checkpoint-marker dicts. Reconstruct the originals before returning.
+        return deserialize_workflow_output(json.loads(metadata.serialized_output))
+
+    async def run_workflow(
+        self,
+        input: Any = None,
+        *,
+        workflow_name: str | None = None,
+        instance_id: str | None = None,
+        wait: bool = True,
+        timeout_seconds: int = 300,
+    ) -> Any:
+        """Start the workflow and, by default, await its output.
+
+        The async counterpart to ``start_workflow`` + ``await_workflow_output``. The
+        underlying durabletask client is synchronous, so the blocking calls run in a
+        worker thread to avoid blocking the event loop.
+
+        Args:
+            input: The initial message/payload for the workflow.
+            workflow_name: The workflow to start. Optional if a default was set on
+                the client; required otherwise.
+            instance_id: Optional explicit orchestration instance ID. If omitted,
+                one is generated.
+            wait: When ``True`` (default), wait for completion and return the
+                deserialized output. When ``False``, return the instance ID as
+                soon as the workflow is scheduled (use with ``stream_workflow`` or
+                the HITL methods).
+            timeout_seconds: Maximum time, in seconds, to wait for completion when
+                ``wait`` is ``True``.
+
+        Returns:
+            The deserialized workflow output when ``wait`` is ``True``; otherwise
+            the orchestration instance ID.
+
+        Raises:
+            TimeoutError: If ``wait`` is ``True`` and the workflow does not complete
+                within ``timeout_seconds``.
+            RuntimeError: If ``wait`` is ``True`` and the workflow ends with a
+                non-successful status.
+        """
+        new_instance_id = await asyncio.to_thread(
+            self.start_workflow, input, workflow_name=workflow_name, instance_id=instance_id
+        )
+        if not wait:
+            return new_instance_id
+        return await asyncio.to_thread(
+            self.await_workflow_output, new_instance_id, workflow_name=workflow_name, timeout_seconds=timeout_seconds
+        )
+
+    def get_runtime_status(self, instance_id: str, *, workflow_name: str | None = None) -> str | None:
+        """Return the workflow's current runtime status name, or ``None`` if unknown.
+
+        Lets callers distinguish a workflow that is still running or paused for
+        human input from one that has reached a terminal state (for example
+        ``COMPLETED``, ``FAILED``, or ``TERMINATED``) — useful when polling, so a
+        workflow that ends without pausing is not mistaken for one that never paused.
+
+        Args:
+            instance_id: The instance ID returned by ``start_workflow``.
+            workflow_name: Optional workflow name; when set (or a client default is
+                set) an instance that does not belong to that workflow returns
+                ``None`` (treated as "not found").
+
+        Returns:
+            The runtime status name (e.g. ``"RUNNING"``, ``"COMPLETED"``), or
+            ``None`` if no state is available for the instance or it belongs to a
+            different workflow.
+        """
+        state = self._client.get_orchestration_state(instance_id)
+        if state is None:
+            return None
+        if not self._is_owned_orchestration(state, workflow_name):
+            return None
+        return state.runtime_status.name
+
+    async def stream_workflow(
+        self,
+        instance_id: str,
+        *,
+        workflow_name: str | None = None,
+        poll_interval_seconds: float = 1.0,
+        timeout_seconds: int | None = None,
+    ) -> AsyncIterator[WorkflowEvent]:
+        """Stream the workflow's events as typed :class:`WorkflowEvent` objects.
+
+        Yields the workflow's events (``executor_invoked`` / ``executor_completed`` /
+        ``output`` / ``request_info`` / ...) in order, finishing when the workflow
+        reaches a terminal state. Each event's ``data`` payload is already
+        reconstructed into its original typed object, so callers do not deserialize
+        anything themselves.
+
+        This is brokerless: it polls the orchestration custom status, into which the
+        orchestrator publishes accumulated events after each superstep. Granularity is
+        per executor and per yielded output, not token-level. Non-agent executors emit
+        events with data payloads; agent executors emit coarse ``executor_invoked`` /
+        ``executor_completed`` lifecycle events. The custom status accumulates events
+        for the run, so this suits workflows with a bounded number of executors rather
+        than very long-running fan-outs.
+
+        Args:
+            instance_id: The instance ID returned by ``start_workflow``.
+            workflow_name: Optional workflow name; when set (or a client default is
+                set) the instance is validated to belong to that workflow before
+                streaming.
+            poll_interval_seconds: Delay between status polls.
+            timeout_seconds: Optional overall timeout; ``None`` streams until the
+                workflow reaches a terminal state.
+
+        Yields:
+            :class:`WorkflowEvent` objects as the workflow progresses.
+
+        Raises:
+            TimeoutError: If ``timeout_seconds`` elapses before completion.
+            ValueError: If the instance does not belong to the targeted workflow.
+        """
+        cursor = 0
+        terminal_statuses = {"COMPLETED", "FAILED", "TERMINATED"}
+        deadline = None if timeout_seconds is None else time.monotonic() + timeout_seconds
+        ownership_checked = False
+
+        while True:
+            state = await asyncio.to_thread(self._client.get_orchestration_state, instance_id)
+
+            # Validate ownership once, on the first poll that returns state.
+            if state is not None and not ownership_checked:
+                if not self._is_owned_orchestration(state, workflow_name):
+                    raise ValueError(f"Instance '{instance_id}' does not belong to the targeted workflow.")
+                ownership_checked = True
+
+            if state is not None:
+                status = self._parse_custom_status(state.serialized_custom_status)
+                if status is not None:
+                    events = status.get("events")
+                    if isinstance(events, list):
+                        typed_events = cast("list[dict[str, Any]]", events)
+                        while cursor < len(typed_events):
+                            yield deserialize_workflow_event(typed_events[cursor])
+                            cursor += 1
+
+            runtime_status = state.runtime_status.name if state is not None else None
+            if runtime_status in terminal_statuses:
+                return
+
+            if deadline is not None and time.monotonic() >= deadline:
+                raise TimeoutError(f"Workflow '{instance_id}' did not complete within {timeout_seconds}s")
+
+            await asyncio.sleep(poll_interval_seconds)
+
+    def get_pending_hitl_requests(self, instance_id: str, *, workflow_name: str | None = None) -> list[dict[str, Any]]:
+        """Return the workflow's pending human-in-the-loop (HITL) requests, if any.
+
+        While a workflow is paused awaiting human input, the orchestrator records the
+        open requests in its custom status. This method reads and normalizes that
+        status so callers do not need to know its internal schema.
+
+        Args:
+            instance_id: The workflow instance ID returned by ``start_workflow``.
+            workflow_name: Optional workflow name; when set (or a client default is
+                set) an instance that does not belong to that workflow returns an
+                empty list (treated as "not found").
+
+        Returns:
+            A list of pending requests. Each entry contains ``request_id``,
+            ``source_executor_id``, ``data``, ``request_type``, and ``response_type``.
+            Empty if the workflow is not currently waiting for human input.
+
+        Note:
+            Requests originating in a nested sub-workflow are included with a
+            **qualified** ``request_id`` (``{executorId}~{ordinal}~{requestId}``, nested
+            for deeper levels). Pass that qualified id straight back to
+            :meth:`send_hitl_response`; it is routed to the owning child orchestration
+            automatically, so the caller only ever addresses the top-level instance.
+        """
+        state = self._client.get_orchestration_state(instance_id)
+        if state is None or not state.serialized_custom_status:
+            return []
+        if not self._is_owned_orchestration(state, workflow_name):
+            return []
+
+        return self._collect_pending_hitl_requests(state.serialized_custom_status)
+
+    @staticmethod
+    def _parse_custom_status(serialized_custom_status: str | None) -> dict[str, Any] | None:
+        """Parse a serialized custom status into a dict, or ``None`` if unusable.
+
+        Returns ``None`` for an empty/absent status or any value that is not a JSON
+        object (the only shape the orchestrator ever writes), so callers can treat
+        "no usable status" uniformly.
+        """
+        if not serialized_custom_status:
+            return None
+        try:
+            parsed = json.loads(serialized_custom_status)
+        except (json.JSONDecodeError, TypeError):
+            return None
+        return cast("dict[str, Any]", parsed) if isinstance(parsed, dict) else None
+
+    def _collect_pending_hitl_requests(self, serialized_custom_status: str) -> list[dict[str, Any]]:
+        """Collect an orchestration's pending requests plus any nested sub-workflow ones.
+
+        Nested requests (discovered via the ``subworkflows`` map the parent records in
+        its custom status as ``{executorId: [childInstanceId, ...]}``) are qualified by
+        ``(executorId, ordinal)`` so deeper requests accumulate a full
+        ``{executorId}~{ordinal}~...~{requestId}`` path and a node with several children
+        keeps each one addressable. Child instances are reached directly by id (already
+        trusted, having come from the parent's status), so no per-child ownership check
+        is applied.
+        """
+        status_dict = self._parse_custom_status(serialized_custom_status)
+        if status_dict is None:
+            return []
+
+        requests: list[dict[str, Any]] = []
+
+        pending = status_dict.get("pending_requests")
+        if isinstance(pending, dict):
+            for request_id, req_data in cast(dict[str, Any], pending).items():
+                if not isinstance(req_data, dict):
+                    continue
+                req = cast(dict[str, Any], req_data)
+                requests.append({
+                    "request_id": req.get("request_id", request_id),
+                    "source_executor_id": req.get("source_executor_id"),
+                    "data": req.get("data"),
+                    "request_type": req.get("request_type"),
+                    "response_type": req.get("response_type"),
+                })
+
+        subworkflows = status_dict.get("subworkflows")
+        if isinstance(subworkflows, dict):
+            for executor_id, child_ids in cast(dict[str, Any], subworkflows).items():
+                children: list[Any] = cast("list[Any]", child_ids) if isinstance(child_ids, list) else []
+                for ordinal, child_instance_id in enumerate(children):
+                    if not isinstance(child_instance_id, str):
+                        continue
+                    child_state = self._client.get_orchestration_state(child_instance_id)
+                    if child_state is None or not child_state.serialized_custom_status:
+                        continue
+                    for child_req in self._collect_pending_hitl_requests(child_state.serialized_custom_status):
+                        qualified = dict(child_req)
+                        qualified["request_id"] = qualify_subworkflow_request_id(
+                            executor_id, ordinal, child_req["request_id"]
+                        )
+                        requests.append(qualified)
+
+        return requests
+
+    def send_hitl_response(
+        self, instance_id: str, request_id: str, response: Any, *, workflow_name: str | None = None
+    ) -> None:
+        """Send a response to a pending HITL request, resuming the workflow.
+
+        The orchestrator correlates the response by using ``request_id`` as the
+        external-event name, so callers do not need to know that convention.
+
+        Args:
+            instance_id: The workflow instance ID.
+            request_id: The pending request's ID (from ``get_pending_hitl_requests``).
+                May be a **qualified** id (``{executorId}~{ordinal}~{requestId}``) for a
+                request that originated in a nested sub-workflow; it is routed to the
+                owning child orchestration automatically.
+            response: The response payload (e.g. a dict matching the expected
+                response type the executor's ``@response_handler`` expects).
+            workflow_name: Optional workflow name; when set (or a client default is
+                set) the instance is validated to belong to that workflow before the
+                event is raised, so a response is never injected into a different
+                workflow's orchestration.
+
+        Raises:
+            ValueError: If the instance does not belong to the targeted workflow, or a
+                qualified id references a sub-workflow that is not currently active.
+
+        Note:
+            The payload is sanitized with ``strip_pickle_markers`` before delivery to
+            neutralize pickle-marker injection, since the worker deserializes it.
+        """
+        # Validate ownership before raising the event when a target is resolvable.
+        if workflow_name or self._default_workflow_name:
+            state = self._client.get_orchestration_state(instance_id)
+            if state is None or not self._is_owned_orchestration(state, workflow_name):
+                raise ValueError(f"Instance '{instance_id}' does not belong to the targeted workflow.")
+
+        # A qualified id addresses a nested sub-workflow: resolve it to the owning child
+        # orchestration instance and the bare request id the child is actually waiting on.
+        target_instance_id, bare_request_id = self._resolve_hitl_target(instance_id, request_id)
+
+        safe_response = strip_pickle_markers(response)
+        self._client.raise_orchestration_event(target_instance_id, event_name=bare_request_id, data=safe_response)
+        logger.debug(
+            "[DurableWorkflowClient] Sent HITL response for request %s on instance %s",
+            bare_request_id,
+            target_instance_id,
+        )
+
+    def _resolve_hitl_target(self, instance_id: str, request_id: str) -> tuple[str, str]:
+        """Resolve a possibly-qualified request id to ``(owning_instance_id, bare_request_id)``.
+
+        An unqualified id (no well-formed hop) targets ``instance_id`` directly. A
+        qualified id ``{executorId}~{ordinal}~{rest}`` addresses a nested sub-workflow:
+        the executor's child instance id is read from this instance's ``subworkflows``
+        custom-status map (a list selected by ``ordinal``) and the remainder is resolved
+        recursively, so arbitrarily deep nesting lands on the leaf child orchestration
+        and its bare request id.
+        """
+        hop = split_subworkflow_request_id(request_id)
+        if hop is None:
+            return instance_id, request_id
+
+        executor_id, ordinal, remainder = hop
+        child_instance_id = self._lookup_subworkflow_instance(instance_id, executor_id, ordinal)
+        if child_instance_id is None:
+            raise ValueError(
+                f"No active sub-workflow '{executor_id}' (ordinal {ordinal}) found for instance "
+                f"'{instance_id}' while routing HITL response for request '{request_id}'."
+            )
+        return self._resolve_hitl_target(child_instance_id, remainder)
+
+    def _lookup_subworkflow_instance(self, instance_id: str, executor_id: str, ordinal: int) -> str | None:
+        """Return the child orchestration instance id for ``(executor_id, ordinal)``, if active.
+
+        Reads the ``subworkflows`` map (``{executorId: [childInstanceId, ...]}``) the
+        parent records in its custom status while dispatching sub-workflow nodes, and
+        selects the child at ``ordinal`` (its dispatch order this superstep).
+        """
+        state = self._client.get_orchestration_state(instance_id)
+        custom_status = self._parse_custom_status(state.serialized_custom_status if state else None)
+        if custom_status is None:
+            return None
+        subworkflows = custom_status.get("subworkflows")
+        if not isinstance(subworkflows, dict):
+            return None
+        children_raw = cast(dict[str, Any], subworkflows).get(executor_id)
+        if not isinstance(children_raw, list):
+            return None
+        children = cast("list[Any]", children_raw)
+        if ordinal < 0 or ordinal >= len(children):
+            return None
+        child = children[ordinal]
+        return child if isinstance(child, str) else None

python/packages/durabletask/agent_framework_durabletask/_workflows/context.py

@@ -0,0 +1,195 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Protocol definition for workflow orchestration contexts.
+
+This module defines the ``WorkflowOrchestrationContext`` protocol that abstracts
+the differences between Azure Functions' ``DurableOrchestrationContext`` and the
+standalone ``durabletask.task.OrchestrationContext``.  The shared workflow
+orchestrator (:func:`run_workflow_orchestrator`) programs against this protocol
+so that the same orchestration logic works on any host.
+
+Each host provides a thin adapter that maps its native context to this protocol:
+
+- ``DurableTaskWorkflowContext`` (this package) — wraps ``OrchestrationContext``
+- ``AzureFunctionsWorkflowContext`` (azurefunctions package) — wraps
+  ``DurableOrchestrationContext``
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class WorkflowOrchestrationContext(Protocol):
+    """Host-agnostic interface for workflow orchestration primitives.
+
+    All methods that return yieldable tasks return ``Any`` because the concrete
+    task types differ between hosting SDKs (``TaskBase`` for Azure Functions,
+    ``Task[T]`` for durabletask).  The generator-based orchestrator simply
+    yields these opaque objects back to the hosting framework.
+    """
+
+    @property
+    def instance_id(self) -> str:
+        """The unique ID of the current orchestration instance."""
+        ...
+
+    @property
+    def is_replaying(self) -> bool:
+        """Whether the orchestrator is replaying previously-recorded history.
+
+        Side effects intended to be observed live exactly once (for example,
+        publishing streaming status to the custom status) must be skipped while
+        this is ``True`` so they are not re-emitted on replay.
+        """
+        ...
+
+    @property
+    def supports_event_streaming(self) -> bool:
+        """Whether this host streams the workflow event timeline via custom status.
+
+        The orchestrator accumulates the full :class:`WorkflowEvent` history and can
+        publish it to the orchestration custom status so a streaming client can
+        replay it (see ``DurableWorkflowClient.stream_workflow``). A host returns
+        ``True`` only when both are true: it has a streaming consumer *and* its
+        custom status can carry an accumulating, payload-bearing event log.
+
+        The Azure Functions host returns ``False``: its Durable Functions custom
+        status is capped at 16 KB (UTF-16) by the WebJobs extension, and its HTTP
+        status endpoint exposes only ``state`` / ``pending_requests`` / ``output``,
+        never the event stream. Publishing the accumulating event log there would
+        overflow the cap and fail the orchestrator without serving any consumer.
+
+        When ``False``, the orchestrator skips event accumulation and omits
+        ``events`` from the custom status; ``state`` and any ``pending_requests``
+        (needed for human-in-the-loop) are still published.
+        """
+        ...
+
+    @property
+    def current_utc_datetime(self) -> datetime:
+        """The current replay-safe UTC datetime."""
+        ...
+
+    def prepare_agent_task(self, executor_id: str, message: str, orchestration_instance_id: str) -> Any:
+        """Create a yieldable task that runs an agent executor.
+
+        Args:
+            executor_id: Agent name / executor ID.
+            message: The text message to send to the agent.
+            orchestration_instance_id: Instance ID used as the entity session key.
+
+        Returns:
+            A yieldable task whose result is an ``AgentResponse``.
+        """
+        ...
+
+    def prepare_activity_task(self, activity_name: str, input_json: str) -> Any:
+        """Create a yieldable task that runs an activity executor.
+
+        Args:
+            activity_name: The registered activity function name.
+            input_json: JSON-serialized activity input.
+
+        Returns:
+            A yieldable task whose result is a JSON string.
+        """
+        ...
+
+    def call_sub_orchestrator(self, name: str, input: Any, instance_id: str | None = None) -> Any:
+        """Create a yieldable task that runs a nested workflow as a child orchestration.
+
+        Used to drive a :class:`~agent_framework.WorkflowExecutor` node: the inner
+        workflow runs as its own durable orchestration (named ``dafx-{innerName}``),
+        independently checkpointed and observable, and its result flows back into
+        the parent's edge routing like any other executor's output.
+
+        Args:
+            name: The registered orchestration name to invoke (``dafx-{innerName}``).
+            input: The JSON-serializable input for the child orchestration.
+            instance_id: Optional deterministic child instance ID. The orchestrator
+                derives one from the parent instance so nested runs are discoverable
+                and replay-safe.
+
+        Returns:
+            A yieldable task whose result is the child orchestration's output.
+        """
+        ...
+
+    def task_all(self, tasks: list[Any]) -> Any:
+        """Create a yieldable composite task that completes when *all* tasks complete.
+
+        Args:
+            tasks: List of yieldable tasks.
+
+        Returns:
+            A yieldable task whose result is a list of individual results.
+        """
+        ...
+
+    def task_any(self, tasks: list[Any]) -> Any:
+        """Create a yieldable composite task that completes when *any* task completes.
+
+        Args:
+            tasks: List of yieldable tasks.
+
+        Returns:
+            A yieldable task whose result is the winning task.
+        """
+        ...
+
+    def wait_for_external_event(self, name: str) -> Any:
+        """Create a yieldable task that waits for a named external event.
+
+        Args:
+            name: Event name to wait for.
+
+        Returns:
+            A yieldable task whose result is the event payload.
+        """
+        ...
+
+    def create_timer(self, fire_at: datetime) -> Any:
+        """Create a yieldable timer task.
+
+        Args:
+            fire_at: UTC datetime when the timer should fire.
+
+        Returns:
+            A yieldable timer task.
+        """
+        ...
+
+    def set_custom_status(self, status: Any) -> None:
+        """Set the orchestration's custom status (visible to external clients).
+
+        Args:
+            status: JSON-serializable status object.
+        """
+        ...
+
+    def new_uuid(self) -> str:
+        """Generate a replay-safe UUID."""
+        ...
+
+    def cancel_task(self, task: Any) -> None:
+        """Best-effort cancellation of a pending task.
+
+        Args:
+            task: The task to cancel.  If the underlying SDK does not support
+                cancellation this is a no-op.
+        """
+        ...
+
+    def get_task_result(self, task: Any) -> Any:
+        """Extract the result from a completed task.
+
+        Args:
+            task: A completed task object.
+
+        Returns:
+            The result value.
+        """
+        ...

python/packages/durabletask/agent_framework_durabletask/_workflows/dt_context.py

@@ -0,0 +1,110 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""DurableTask SDK adapter for WorkflowOrchestrationContext.
+
+Wraps ``durabletask.task.OrchestrationContext`` to satisfy the
+:class:`WorkflowOrchestrationContext` protocol.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+from typing import Any, cast
+
+from durabletask.task import (
+    OrchestrationContext,
+    Task,
+    when_all,
+    when_any,
+)
+
+from .._executors import OrchestrationAgentExecutor
+from .._models import AgentSessionId, DurableAgentSession
+from .._shim import DurableAIAgent
+from .context import WorkflowOrchestrationContext
+
+logger = logging.getLogger(__name__)
+
+
+class DurableTaskWorkflowContext:
+    """Adapter that maps ``OrchestrationContext`` to :class:`WorkflowOrchestrationContext`."""
+
+    def __init__(self, context: OrchestrationContext) -> None:
+        self._context = context
+        self._executor = OrchestrationAgentExecutor(context)
+
+    # -- Properties -----------------------------------------------------------
+
+    @property
+    def instance_id(self) -> str:
+        return self._context.instance_id
+
+    @property
+    def is_replaying(self) -> bool:
+        return self._context.is_replaying
+
+    @property
+    def supports_event_streaming(self) -> bool:
+        # The standalone DurableTask host exposes the event timeline to clients via
+        # DurableWorkflowClient.stream_workflow, and its DTS backend imposes no 16 KB
+        # custom-status cap, so the full accumulated event stream is published.
+        return True
+
+    @property
+    def current_utc_datetime(self) -> datetime:
+        return self._context.current_utc_datetime
+
+    # -- Agent / Activity dispatch --------------------------------------------
+
+    def prepare_agent_task(self, executor_id: str, message: str, orchestration_instance_id: str) -> Any:
+        session_id = AgentSessionId(name=executor_id, key=orchestration_instance_id)
+        session = DurableAgentSession(durable_session_id=session_id)
+        agent = DurableAIAgent(self._executor, executor_id)
+        return agent.run(message, session=session)
+
+    def prepare_activity_task(self, activity_name: str, input_json: str) -> Any:
+        return cast(Any, self._context.call_activity(activity_name, input=input_json))
+
+    def call_sub_orchestrator(self, name: str, input: Any, instance_id: str | None = None) -> Any:
+        return cast(Any, self._context.call_sub_orchestrator(name, input=input, instance_id=instance_id))
+
+    # -- Composite tasks ------------------------------------------------------
+
+    def task_all(self, tasks: list[Any]) -> Any:
+        return when_all(tasks)
+
+    def task_any(self, tasks: list[Any]) -> Any:
+        return when_any(tasks)
+
+    # -- External events / timers ---------------------------------------------
+
+    def wait_for_external_event(self, name: str) -> Any:
+        return cast(Any, self._context).wait_for_external_event(name)
+
+    def create_timer(self, fire_at: datetime) -> Any:
+        return cast(Any, self._context).create_timer(fire_at)
+
+    # -- Status / utility -----------------------------------------------------
+
+    def set_custom_status(self, status: Any) -> None:
+        self._context.set_custom_status(status)
+
+    def new_uuid(self) -> str:
+        return self._context.new_uuid()
+
+    def cancel_task(self, task: Any) -> None:
+        # durabletask Task doesn't expose cancel(); this is a best-effort no-op.
+        cancel_fn = getattr(task, "cancel", None)
+        if callable(cancel_fn):
+            cancel_fn()
+
+    def get_task_result(self, task: Any) -> Any:
+        if isinstance(task, Task):
+            return cast(Any, task.get_result())
+        return getattr(task, "result", None)
+
+
+# Ensure the adapter satisfies the protocol. Validated statically by the type
+# checker (and at every ``run_workflow_orchestrator`` call site) with no runtime cost.
+_protocol_check: type[WorkflowOrchestrationContext] = DurableTaskWorkflowContext

python/packages/durabletask/agent_framework_durabletask/_workflows/naming.py

@@ -0,0 +1,299 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Durable naming helpers for hosting MAF Workflows.
+
+A hosted workflow maps to durable primitives (an orchestration, plus an activity
+or entity per executor) whose names must be **stable** across worker restarts:
+durable replay only resumes an in-flight orchestration if the orchestration,
+activity, and entity names still resolve to the same functions. This module
+centralizes how those names are derived from a workflow name so every host (the
+Azure Functions host and the standalone durabletask worker) and the client agree
+on one scheme.
+
+Naming scheme (the orchestration name is aligned byte-for-byte with .NET's
+``WorkflowNamingHelper``)::
+
+    orchestration:       dafx-{workflowName}
+    non-agent activity:  dafx-{workflowName}-{executorId}
+    agent entity:        dafx-{workflowName}-{executorId}
+
+The orchestration name is the identifier the Durable Task tooling/UI surfaces, so
+it matches .NET exactly. The inner activity/entity names are scoped by workflow in
+Python (unlike .NET's bare ``dafx-{executorId}``) so two co-hosted workflows that
+reuse an executor id cannot collide.
+"""
+
+from __future__ import annotations
+
+import re
+
+__all__ = [
+    "DURABLE_NAME_PREFIX",
+    "MAX_EXECUTOR_ID_LENGTH",
+    "SUBWORKFLOW_REQUEST_SEPARATOR",
+    "is_auto_generated_workflow_name",
+    "qualify_subworkflow_request_id",
+    "split_subworkflow_request_id",
+    "validate_executor_id",
+    "validate_workflow_name",
+    "workflow_executor_activity_name",
+    "workflow_name_from_orchestrator",
+    "workflow_orchestrator_name",
+    "workflow_scoped_executor_id",
+]
+
+# Shared prefix for every durable name this hosting layer registers. Matches
+# .NET's ``WorkflowNamingHelper.OrchestrationFunctionPrefix`` and the existing
+# ``AgentSessionId.ENTITY_NAME_PREFIX``.
+DURABLE_NAME_PREFIX = "dafx-"
+
+# Separator used to qualify a nested sub-workflow's pending HITL request when it is
+# bubbled up to the top-level instance (one top-level addressing surface). A qualified id
+# is a path of ``{executorId}~{ordinal}`` hops ending in the leaf's bare request id,
+# e.g. ``review~0~approve~1~<requestId>``. Both hosts and the client must agree on it
+# so a qualified id round-trips: the read side prepends hops; the respond side peels
+# them to route the response to the owning child orchestration.
+#
+# ``~`` (RFC 3986 "unreserved", so URL-path-safe) is deliberately **not** ``::``:
+# core emits ``auto::{index}`` request ids for functional ``@workflow`` HITL, so a
+# ``::`` separator would mis-parse those leaf ids. ``~`` does not appear in core
+# request ids (uuid4 or ``auto::N``); executor ids are validated to exclude it (see
+# :func:`validate_executor_id`), so only the structural hops carry the separator.
+SUBWORKFLOW_REQUEST_SEPARATOR = "~"
+
+# Upper bound on an executor id's length when a workflow is hosted durably. The id is
+# interpolated into durable activity/entity names (``dafx-{workflow}-{executor}``) and,
+# for sub-workflow nodes, into recursively-nested child orchestration instance ids
+# (``{parent}::{executor}::{n}``). Capping it keeps those derived strings within typical
+# durable backend name/id limits; combined with the workflow-name cap, the worst-case
+# instance id stays bounded even for deeply-nested sub-workflows.
+MAX_EXECUTOR_ID_LENGTH = 128
+
+# A workflow name is interpolated into durable orchestration/activity/entity names
+# *and* into HTTP route segments (``workflow/{workflowName}/run``), so it must be
+# conservative enough to be safe in every position: ASCII letters, digits, '_' or
+# '-', starting with a letter, at most 63 characters. The length cap leaves room
+# for the ``dafx-`` prefix and an ``-{executorId}`` suffix within typical durable
+# name limits.
+_WORKFLOW_NAME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_-]{0,62}$")
+
+# Names auto-generated by ``WorkflowBuilder`` when the caller does not pass one,
+# e.g. ``"WorkflowBuilder-3f2b1c0a-1234-5678-9abc-def012345678"``. They embed a
+# fresh ``uuid4`` per process build, so they are not stable identities and must be
+# rejected for durable hosting (see :func:`validate_workflow_name`).
+_AUTO_GENERATED_NAME_RE = re.compile(
+    r"^WorkflowBuilder-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$"
+)
+
+
+def workflow_orchestrator_name(workflow_name: str) -> str:
+    """Return the durable orchestration name for a workflow.
+
+    Args:
+        workflow_name: The workflow's name. Must satisfy
+            :func:`validate_workflow_name`.
+
+    Returns:
+        ``"dafx-{workflow_name}"``.
+
+    Raises:
+        ValueError: If ``workflow_name`` is not a valid, stable workflow name.
+    """
+    validate_workflow_name(workflow_name)
+    return f"{DURABLE_NAME_PREFIX}{workflow_name}"
+
+
+def workflow_name_from_orchestrator(orchestrator_name: str) -> str | None:
+    """Recover the workflow name from a durable orchestration name.
+
+    The inverse of :func:`workflow_orchestrator_name`. Intended to be applied to
+    orchestration names (for example a durable instance's ``status.name``); it
+    strips the shared :data:`DURABLE_NAME_PREFIX`.
+
+    Args:
+        orchestrator_name: A durable orchestration name.
+
+    Returns:
+        The workflow name, or ``None`` if ``orchestrator_name`` does not carry the
+        expected prefix (so a caller can treat it as "not one of ours").
+    """
+    if not orchestrator_name.startswith(DURABLE_NAME_PREFIX):
+        return None
+    name = orchestrator_name[len(DURABLE_NAME_PREFIX) :]
+    return name or None
+
+
+def workflow_scoped_executor_id(workflow_name: str, executor_id: str) -> str:
+    """Return the workflow-scoped identity for an executor.
+
+    Inner executors (non-agent activities and agent entities) are scoped by
+    workflow so two co-hosted workflows that reuse an ``executor_id`` register and
+    dispatch to distinct durable primitives instead of colliding on one global
+    name. This is the **unprefixed** identity (e.g. used as
+    :class:`~agent_framework_durabletask.AgentSessionId` ``name``, which the entity
+    layer then prefixes); see :func:`workflow_executor_activity_name` for the full
+    activity function name.
+
+    Args:
+        workflow_name: The owning workflow's name.
+        executor_id: The executor's id within that workflow.
+
+    Returns:
+        ``"{workflow_name}-{executor_id}"``.
+    """
+    return f"{workflow_name}-{executor_id}"
+
+
+def workflow_executor_activity_name(workflow_name: str, executor_id: str) -> str:
+    """Return the durable activity function name for a non-agent executor.
+
+    Args:
+        workflow_name: The owning workflow's name.
+        executor_id: The executor's id within that workflow.
+
+    Returns:
+        ``"dafx-{workflow_name}-{executor_id}"``.
+    """
+    return f"{DURABLE_NAME_PREFIX}{workflow_scoped_executor_id(workflow_name, executor_id)}"
+
+
+def validate_workflow_name(workflow_name: str) -> None:
+    """Validate that a workflow name is usable as a stable durable identity.
+
+    The name is **validated and rejected** rather than silently sanitized. A
+    workflow name is an identity baked into durable orchestration/activity/entity
+    names and HTTP routes, so transforming it could either (a) collapse two
+    distinct names into one and reintroduce the cross-workflow collision this
+    scheme exists to prevent, or (b) change the resolved name across versions and
+    break resume of in-flight instances. A loud error is safer than a silent
+    rename.
+
+    Args:
+        workflow_name: The candidate name.
+
+    Raises:
+        ValueError: If the name is empty, an auto-generated ``WorkflowBuilder``
+            name, or contains characters outside
+            ``[A-Za-z][A-Za-z0-9_-]{0,62}``.
+    """
+    if not workflow_name:
+        raise ValueError("Workflow name must be a non-empty string.")
+    if is_auto_generated_workflow_name(workflow_name):
+        raise ValueError(
+            f"Workflow name '{workflow_name}' is an auto-generated WorkflowBuilder name, which is "
+            "not stable across restarts. Pass an explicit, stable name to WorkflowBuilder(name=...) "
+            "before hosting the workflow durably."
+        )
+    if not _WORKFLOW_NAME_RE.match(workflow_name):
+        raise ValueError(
+            f"Workflow name '{workflow_name}' is invalid. Use 1-63 characters consisting of ASCII "
+            "letters, digits, '_' or '-', and starting with a letter."
+        )
+
+
+def is_auto_generated_workflow_name(workflow_name: str) -> bool:
+    """Return whether a name looks like ``WorkflowBuilder``'s auto-generated default.
+
+    ``WorkflowBuilder`` names an otherwise-unnamed workflow
+    ``f"WorkflowBuilder-{uuid4()}"``, which changes on every process build and is
+    therefore not a stable durable identity.
+
+    Args:
+        workflow_name: The candidate name.
+
+    Returns:
+        ``True`` if the name matches the auto-generated pattern.
+    """
+    return bool(_AUTO_GENERATED_NAME_RE.match(workflow_name))
+
+
+def validate_executor_id(executor_id: str) -> None:
+    """Validate that an executor id is safe to host durably.
+
+    An executor id is interpolated into durable activity/entity names and, for
+    sub-workflow nodes, into nested child-orchestration instance ids and the
+    qualified ids used to address nested human-in-the-loop requests. Two properties
+    must hold:
+
+    * It must not contain :data:`SUBWORKFLOW_REQUEST_SEPARATOR`. That sequence
+      separates the structural hops of a qualified nested-HITL request id, so an id
+      containing it would make a qualified id ambiguous and mis-route a response.
+    * It must be at most :data:`MAX_EXECUTOR_ID_LENGTH` characters, so the durable
+      names and (recursively nested) instance ids derived from it stay within typical
+      durable backend limits.
+
+    Args:
+        executor_id: The executor's id within a hosted workflow.
+
+    Raises:
+        ValueError: If the id is empty, contains the reserved separator, or is too
+            long.
+    """
+    if not executor_id:
+        raise ValueError("Executor id must be a non-empty string.")
+    if SUBWORKFLOW_REQUEST_SEPARATOR in executor_id:
+        raise ValueError(
+            f"Executor id '{executor_id}' contains the reserved sub-workflow request separator "
+            f"'{SUBWORKFLOW_REQUEST_SEPARATOR}', which is used to address nested human-in-the-loop "
+            "requests. Rename the executor so its id does not contain that sequence."
+        )
+    if len(executor_id) > MAX_EXECUTOR_ID_LENGTH:
+        raise ValueError(
+            f"Executor id '{executor_id[:32]}...' is too long ({len(executor_id)} > "
+            f"{MAX_EXECUTOR_ID_LENGTH}). Durable activity/entity names and nested instance ids are "
+            "derived from it; use a shorter id."
+        )
+
+
+def qualify_subworkflow_request_id(executor_id: str, ordinal: int, inner_request_id: str) -> str:
+    """Prepend one sub-workflow hop to a (possibly already-qualified) request id.
+
+    Produces ``{executor_id}~{ordinal}~{inner_request_id}``. ``ordinal`` selects the
+    specific child orchestration among several a single ``WorkflowExecutor`` node may
+    dispatch in one superstep, so two children of the same executor stay distinctly
+    addressable. ``inner_request_id`` is the child's bare leaf request id or its own
+    already-qualified path for deeper nesting.
+
+    Args:
+        executor_id: The sub-workflow node's executor id (separator-free; see
+            :func:`validate_executor_id`).
+        ordinal: The child's index in the parent's ``subworkflows`` status list.
+        inner_request_id: The request id (bare or qualified) within the child.
+
+    Returns:
+        The qualified request id one level higher.
+    """
+    sep = SUBWORKFLOW_REQUEST_SEPARATOR
+    return f"{executor_id}{sep}{ordinal}{sep}{inner_request_id}"
+
+
+def split_subworkflow_request_id(request_id: str) -> tuple[str, int, str] | None:
+    """Peel the outermost sub-workflow hop off a qualified request id.
+
+    The inverse of :func:`qualify_subworkflow_request_id` for a single level.
+    Returns ``(executor_id, ordinal, remainder)`` where ``remainder`` is the still
+    (possibly) qualified id one level deeper, or ``None`` when ``request_id`` carries
+    no well-formed hop -- i.e. it is a bare leaf request id that targets the current
+    instance directly. A leaf id may itself contain the separator (e.g. core's
+    ``auto::N`` does not, but a custom id could); because only structural hops use the
+    ``{executor}~{int-ordinal}~`` shape, a value whose second segment is not an integer
+    is treated as a bare leaf rather than a hop.
+
+    Args:
+        request_id: A bare or qualified request id.
+
+    Returns:
+        ``(executor_id, ordinal, remainder)`` for a qualified id, else ``None``.
+    """
+    sep = SUBWORKFLOW_REQUEST_SEPARATOR
+    if sep not in request_id:
+        return None
+    parts = request_id.split(sep, 2)
+    if len(parts) < 3:
+        return None
+    executor_id, ordinal_str, remainder = parts
+    try:
+        ordinal = int(ordinal_str)
+    except ValueError:
+        return None
+    return executor_id, ordinal, remainder

python/packages/durabletask/agent_framework_durabletask/_workflows/registration.py

@@ -0,0 +1,137 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Host-agnostic plan for registering a MAF Workflow as a durable orchestration.
+
+A MAF :class:`Workflow` is hosted by turning each graph node into a durable
+primitive:
+
+- each :class:`AgentExecutor` becomes a durable **entity**,
+- each :class:`WorkflowExecutor` (a nested sub-workflow) becomes a durable
+  **child orchestration**, and
+- each other :class:`Executor` becomes a durable **activity**,
+
+driven by a single workflow **orchestrator**.
+
+The *decision* of which executor maps to which primitive is identical on every
+host (Azure Functions or a standalone durabletask worker); only the *mechanism*
+for registering them differs (Functions trigger decorators vs.
+``worker.add_*``). :func:`plan_workflow_registration` captures the shared
+decision so each host applies one consistent plan with its own registration
+mechanism — analogous to .NET's shared ``DurableWorkflowOptions`` feeding
+host-specific trigger generation.
+
+Sub-workflows nest: a hosted workflow may contain :class:`WorkflowExecutor`
+nodes whose inner workflows must themselves be registered (their orchestrator,
+agents, and activities) so the parent can drive them via
+``call_sub_orchestrator``. :func:`collect_hosted_workflows` walks that tree so a
+host registers every reachable workflow exactly once.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+from agent_framework import AgentExecutor, Executor, Workflow, WorkflowExecutor
+
+
+@dataclass
+class WorkflowRegistrationPlan:
+    """The durable primitives a workflow registers, independent of host.
+
+    Attributes:
+        agent_executors: Agent executors to register as durable entities. The
+            full :class:`AgentExecutor` is carried (not just its agent) so each
+            host can register the entity under the executor's ``id`` — the same
+            identity the orchestrator dispatches to — which keeps
+            ``AgentExecutor(agent, id=...)`` working when the id differs from
+            ``agent.name``.
+        activity_executors: Non-agent, non-subworkflow executors to register as
+            durable activities.
+        subworkflow_executors: :class:`WorkflowExecutor` nodes whose inner
+            workflows are driven as durable child orchestrations. The node itself
+            is *not* registered as an activity; its inner workflow is registered
+            separately (see :func:`collect_hosted_workflows`).
+    """
+
+    agent_executors: list[AgentExecutor]
+    activity_executors: list[Executor]
+    subworkflow_executors: list[WorkflowExecutor]
+
+
+def plan_workflow_registration(workflow: Workflow) -> WorkflowRegistrationPlan:
+    """Classify a workflow's executors into the durable primitives to register.
+
+    Args:
+        workflow: The MAF :class:`Workflow` to host.
+
+    Returns:
+        A :class:`WorkflowRegistrationPlan` describing the agent executors
+        (entities), sub-workflow executors (child orchestrations), and the
+        remaining non-agent executors (activities).
+    """
+    agent_executors: list[AgentExecutor] = []
+    activity_executors: list[Executor] = []
+    subworkflow_executors: list[WorkflowExecutor] = []
+
+    for executor in workflow.executors.values():
+        if isinstance(executor, AgentExecutor):
+            agent_executors.append(executor)
+        elif isinstance(executor, WorkflowExecutor):
+            subworkflow_executors.append(executor)
+        else:
+            activity_executors.append(executor)
+
+    return WorkflowRegistrationPlan(
+        agent_executors=agent_executors,
+        activity_executors=activity_executors,
+        subworkflow_executors=subworkflow_executors,
+    )
+
+
+def collect_hosted_workflows(workflow: Workflow) -> Iterator[Workflow]:
+    """Yield ``workflow`` and every nested sub-workflow, deduped by name.
+
+    A host registers the orchestration primitives for each yielded workflow so a
+    parent orchestration can invoke its sub-workflows as child orchestrations.
+    Workflows are deduped by :attr:`Workflow.name`, **compared case-insensitively**:
+    the *same* sub-workflow instance reused across the tree (or shared by two
+    top-level workflows) is yielded once, which is the expected fan-out pattern. Two
+    **different** workflow instances whose names collide (including case-only
+    differences) are rejected, since both would resolve to one durable orchestration
+    (``dafx-{name}``) -- whose name the route ownership check compares
+    case-insensitively -- and would silently shadow each other. The top-level
+    ``workflow`` is yielded first.
+
+    Args:
+        workflow: The top-level workflow to walk.
+
+    Yields:
+        Each distinct workflow in the nesting tree, parent before child.
+
+    Raises:
+        ValueError: If two different workflow instances in the tree have colliding
+            (case-insensitive) names.
+    """
+    seen: dict[str, Workflow] = {}
+
+    def _walk(current: Workflow) -> Iterator[Workflow]:
+        key = current.name.casefold()
+        existing = seen.get(key)
+        if existing is not None:
+            if existing is not current:
+                raise ValueError(
+                    f"A different workflow named '{current.name}' collides with '{existing.name}'. A "
+                    f"workflow name maps to a single durable orchestration ('dafx-{current.name}'), "
+                    "compared case-insensitively, so names must be unique within a hosted composition. "
+                    "Rename one, or reuse the same Workflow instance if they are meant to be the same "
+                    "sub-workflow."
+                )
+            return
+        seen[key] = current
+        yield current
+        plan = plan_workflow_registration(current)
+        for sub in plan.subworkflow_executors:
+            yield from _walk(sub.workflow)
+
+    yield from _walk(workflow)

python/packages/durabletask/agent_framework_durabletask/_workflows/runner_context.py

@@ -0,0 +1,147 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Runner context for activity execution within durable orchestrations.
+
+This module provides the :class:`CapturingRunnerContext` class that captures
+messages and events produced during executor execution within activities.
+It is host-agnostic and works on any durable task host.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from copy import copy
+from typing import Any
+
+from agent_framework import (
+    CheckpointStorage,
+    RunnerContext,
+    WorkflowCheckpoint,
+    WorkflowEvent,
+    WorkflowMessage,
+)
+from agent_framework._workflows._runner_context import YieldOutputClassifier, YieldOutputEventType
+from agent_framework._workflows._state import State
+
+
+class CapturingRunnerContext(RunnerContext):
+    """A RunnerContext that captures messages and events for durable activities.
+
+    This context captures all messages and events produced during execution
+    without requiring durable entity storage, allowing the results to be
+    returned to the orchestrator.
+
+    Checkpointing is not supported — the orchestrator manages state.
+    """
+
+    def __init__(self) -> None:
+        self._messages: dict[str, list[WorkflowMessage]] = {}
+        self._event_queue: asyncio.Queue[WorkflowEvent] = asyncio.Queue()
+        self._pending_request_info_events: dict[str, WorkflowEvent[Any]] = {}
+        self._workflow_id: str | None = None
+        self._streaming: bool = False
+        self._yield_output_classifier: YieldOutputClassifier = lambda _executor_id: "output"
+
+    # -- Messaging ------------------------------------------------------------
+
+    async def send_message(self, message: WorkflowMessage) -> None:
+        self._messages.setdefault(message.source_id, [])
+        self._messages[message.source_id].append(message)
+
+    async def drain_messages(self) -> dict[str, list[WorkflowMessage]]:
+        messages = copy(self._messages)
+        self._messages.clear()
+        return messages
+
+    async def has_messages(self) -> bool:
+        return bool(self._messages)
+
+    # -- Events ---------------------------------------------------------------
+
+    async def add_event(self, event: WorkflowEvent) -> None:
+        await self._event_queue.put(event)
+
+    async def drain_events(self) -> list[WorkflowEvent]:
+        events: list[WorkflowEvent] = []
+        while True:
+            try:
+                events.append(self._event_queue.get_nowait())
+            except asyncio.QueueEmpty:
+                break
+        return events
+
+    async def has_events(self) -> bool:
+        return not self._event_queue.empty()
+
+    async def next_event(self) -> WorkflowEvent:
+        return await self._event_queue.get()
+
+    # -- Checkpointing (not supported) ----------------------------------------
+
+    def has_checkpointing(self) -> bool:
+        return False
+
+    def set_runtime_checkpoint_storage(self, storage: CheckpointStorage) -> None:
+        pass
+
+    def clear_runtime_checkpoint_storage(self) -> None:
+        pass
+
+    async def create_checkpoint(
+        self,
+        workflow_name: str,
+        graph_signature_hash: str,
+        state: State,
+        previous_checkpoint_id: str | None,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:
+        raise NotImplementedError("Checkpointing is not supported in activity context")
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        raise NotImplementedError("Checkpointing is not supported in activity context")
+
+    async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
+        raise NotImplementedError("Checkpointing is not supported in activity context")
+
+    # -- Workflow configuration -----------------------------------------------
+
+    def set_workflow_id(self, workflow_id: str) -> None:
+        self._workflow_id = workflow_id
+
+    def reset_for_new_run(self) -> None:
+        self._messages.clear()
+        self._event_queue = asyncio.Queue()
+        self._pending_request_info_events.clear()
+        self._streaming = False
+
+    def set_streaming(self, streaming: bool) -> None:
+        self._streaming = streaming
+
+    def is_streaming(self) -> bool:
+        return self._streaming
+
+    # -- Yield-output classification -------------------------------------------
+
+    def set_yield_output_classifier(self, classifier: YieldOutputClassifier) -> None:
+        """Set the classifier used by ``WorkflowContext.yield_output()``."""
+        self._yield_output_classifier = classifier
+
+    def classify_yielded_output(self, executor_id: str) -> YieldOutputEventType | None:
+        """Classify an executor's yield_output payload as output, intermediate, or hidden."""
+        return self._yield_output_classifier(executor_id)
+
+    # -- Request Info Events --------------------------------------------------
+
+    async def add_request_info_event(self, event: WorkflowEvent[Any]) -> None:
+        self._pending_request_info_events[event.request_id] = event
+        await self.add_event(event)
+
+    async def send_request_info_response(self, request_id: str, response: Any) -> None:
+        raise NotImplementedError(
+            "send_request_info_response is not supported in activity context. "
+            "Human-in-the-loop scenarios should be handled at the orchestrator level."
+        )
+
+    async def get_pending_request_info_events(self) -> dict[str, WorkflowEvent[Any]]:
+        return dict(self._pending_request_info_events)

python/packages/durabletask/agent_framework_durabletask/_workflows/serialization.py

@@ -0,0 +1,357 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Internal serialization helpers for workflow execution.
+
+These helpers are framework-internal plumbing for moving typed objects across
+durable orchestration/activity boundaries. They are **not** part of the public
+API and must not be called by application code.
+
+They wrap the core checkpoint codec (``encode_checkpoint_value`` /
+``decode_checkpoint_value`` from ``agent_framework._workflows``), which uses
+pickle + base64 to round-trip arbitrary Python objects (dataclasses, Pydantic
+models, ``Message``, ...) while leaving JSON-native types (str, int, float,
+bool, None) as-is.
+
+Because that codec can unpickle objects, every value that crosses an external
+trust boundary -- HTTP request bodies and HITL responses raised as external
+events -- is sanitized by the framework with :func:`strip_pickle_markers`
+*before* it can reach these helpers. Application code never has to perform that
+sanitization itself: the orchestrator, the activity body, and the HTTP entry
+points already do it at the boundary. See
+:mod:`agent_framework._workflows._checkpoint_encoding` for the full security model.
+
+Contents:
+- ``serialize_value`` / ``deserialize_value``: internal codec aliases for encode/decode.
+- ``reconstruct_to_type``: rebuilds HITL response data (which arrives without type
+  markers) to a known type.
+- ``resolve_type``: resolves 'module:class' type keys to Python types.
+- ``strip_pickle_markers``: the framework's trust-boundary defense that neutralizes
+  attacker-injected pickle/type markers.
+"""
+
+from __future__ import annotations
+
+import importlib
+import logging
+from contextlib import suppress
+from dataclasses import is_dataclass
+from typing import Any, cast
+
+from agent_framework import WorkflowEvent
+from agent_framework._workflows._checkpoint_encoding import (
+    _PICKLE_MARKER,  # pyright: ignore[reportPrivateUsage]
+    _TYPE_MARKER,  # pyright: ignore[reportPrivateUsage]
+    decode_checkpoint_value,
+    encode_checkpoint_value,
+)
+from agent_framework._workflows._events import WorkflowEventType
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_type(type_key: str) -> type | None:
+    """Resolve a 'module:class' type key to its Python type.
+
+    Args:
+        type_key: Fully qualified type reference in 'module_name:class_name' format.
+
+    Returns:
+        The resolved type, or None if resolution fails.
+    """
+    try:
+        module_name, class_name = type_key.split(":", 1)
+        module = importlib.import_module(module_name)
+        resolved = getattr(module, class_name, None)
+        # Only return actual classes. A non-type attribute (function, module member,
+        # etc.) would raise TypeError in issubclass() inside reconstruct_to_type().
+        return resolved if isinstance(resolved, type) else None
+    except Exception:
+        logger.debug("Could not resolve type %s", type_key)
+        return None
+
+
+# ============================================================================
+# Pickle marker sanitization (security)
+# ============================================================================
+
+
+def strip_pickle_markers(data: Any) -> Any:
+    """Recursively strip pickle/type markers from untrusted data.
+
+    The core checkpoint encoding uses ``__pickled__`` and ``__type__`` markers to
+    roundtrip arbitrary Python objects via *pickle*.  If an attacker crafts an
+    HTTP payload that contains these markers, the data would flow into
+    ``pickle.loads()`` and enable **arbitrary code execution**.
+
+    This function walks the incoming data structure and replaces any ``dict``
+    that contains either marker key with ``None``, neutralizing the attack
+    vector while leaving all other data untouched.
+
+    The framework applies this at every external trust boundary -- HTTP request
+    bodies and HITL responses raised as external events -- before the value can
+    reach the internal codec (:func:`deserialize_value` /
+    ``decode_checkpoint_value``). Application code does not need to call it.
+    """
+    if isinstance(data, dict):
+        if _PICKLE_MARKER in data or _TYPE_MARKER in data:
+            logger.debug("Stripped pickle/type markers from untrusted input.")
+            return None
+        typed_dict = cast(dict[str, Any], data)
+        return {k: strip_pickle_markers(v) for k, v in typed_dict.items()}
+
+    if isinstance(data, list):
+        typed_list = cast(list[Any], data)
+        return [strip_pickle_markers(item) for item in typed_list]
+
+    return data
+
+
+# ============================================================================
+# Sub-workflow envelope markers (trust boundary)
+# ============================================================================
+
+# A WorkflowExecutor node runs its inner workflow as a durable child orchestration.
+# The parent wraps the node's input in this envelope so the child orchestrator can
+# tell a trusted sub-orchestration payload (serialized by the parent, post-boundary,
+# via call_sub_orchestrator) apart from untrusted top-level client input.
+SUBWORKFLOW_INPUT_KEY = "__subworkflow_input__"
+
+# When a workflow runs as a sub-workflow, its orchestrator returns this envelope
+# instead of a bare outputs list, so the parent can recover both the inner outputs
+# *and* the inner event timeline (a child orchestration is a separate durable
+# instance; its return value is the only deterministic, replay-safe channel back to
+# the parent). A top-level run still returns a bare list, so the client output path
+# is unchanged. See ``orchestrator._process_subworkflow_result``.
+SUBWORKFLOW_RESULT_KEY = "__subworkflow_result__"
+
+
+def strip_subworkflow_markers(data: Any) -> Any:
+    """Remove the reserved sub-workflow envelope key from untrusted top-level input.
+
+    The orchestrator treats a top-level input dict carrying :data:`SUBWORKFLOW_INPUT_KEY`
+    as a *trusted* child-orchestration payload and reconstructs it with
+    :func:`deserialize_value` (pickle) **without** the usual
+    :func:`strip_pickle_markers` sanitization, because a genuine envelope is only ever
+    built internally (post trust boundary) by ``call_sub_orchestrator``. If untrusted
+    client input could carry that key, an attacker could smuggle a pickle payload
+    straight into ``pickle.loads`` (RCE).
+
+    Hosts therefore call this on client-supplied workflow input *before* scheduling the
+    orchestration, so the only way the orchestrator ever sees the envelope is from a
+    real internal child dispatch. Only the top-level key is removed (that is the only
+    position the orchestrator interprets it), leaving the rest of the caller's payload
+    untouched.
+    """
+    if not isinstance(data, dict):
+        return data
+    typed = cast(dict[str, Any], data)
+    if SUBWORKFLOW_INPUT_KEY not in typed:
+        return typed
+    logger.debug("Stripped reserved sub-workflow envelope key from untrusted input.")
+    cleaned = typed.copy()
+    cleaned.pop(SUBWORKFLOW_INPUT_KEY, None)
+    return cleaned
+
+
+# ============================================================================
+# Serialize / Deserialize
+# ============================================================================
+
+
+def serialize_value(value: Any) -> Any:
+    """Encode a value for JSON-compatible cross-activity communication (internal).
+
+    Framework-internal codec. Delegates to core checkpoint encoding which uses
+    pickle + base64 for non-JSON-native types (dataclasses, Pydantic models,
+    Message, etc.). Not part of the public API.
+
+    Args:
+        value: Any Python value (primitive, dataclass, Pydantic model, Message, etc.)
+
+    Returns:
+        A JSON-serializable representation with embedded type metadata for reconstruction.
+    """
+    return encode_checkpoint_value(value)
+
+
+def deserialize_value(value: Any) -> Any:
+    """Decode a value previously encoded with :func:`serialize_value` (internal).
+
+    Framework-internal codec. Delegates to core checkpoint decoding which
+    unpickles base64-encoded values and verifies type integrity. Not part of the
+    public API: callers only ever hand it values that the framework produced
+    itself or that have already passed the :func:`strip_pickle_markers` trust
+    boundary, so untrusted markers can never reach ``pickle.loads()`` here.
+
+    Args:
+        value: The serialized data (dict with pickle markers, list, or primitive)
+
+    Returns:
+        Reconstructed typed object if type metadata found, otherwise original value.
+    """
+    return decode_checkpoint_value(value)
+
+
+def deserialize_workflow_output(output: Any) -> Any:
+    """Reconstruct the workflow outputs produced by the shared activity.
+
+    Each value an executor yields is encoded with :func:`serialize_value` before
+    it reaches the orchestrator, so typed objects (dataclasses, Pydantic models,
+    ``AgentResponse``, ...) are stored as checkpoint-marker dicts. This reverses
+    that encoding so callers receive the original objects.
+
+    This is the single decode path shared by every host (the in-process
+    :class:`DurableWorkflowClient` and the Azure Functions status endpoint) so
+    they never diverge in how a completed workflow's output is reconstructed.
+
+    ``output`` must originate from the workflow's own orchestration result
+    (trusted durable storage), never from untrusted external input. Markers in
+    untrusted input must be neutralized with :func:`strip_pickle_markers` first.
+
+    Args:
+        output: The workflow's orchestration result, already JSON-decoded (a list
+            of yielded outputs or a single value).
+
+    Returns:
+        The output with every checkpoint-encoded value reconstructed; primitives
+        and plain JSON structures pass through unchanged.
+    """
+    return deserialize_value(output)
+
+
+# ============================================================================
+# Workflow Event Serialization (streaming)
+# ============================================================================
+
+
+def _type_key(value_type: type[Any] | None) -> str | None:
+    """Format a type as a ``'module:qualname'`` key for :func:`resolve_type`."""
+    if value_type is None:
+        return None
+    return f"{value_type.__module__}:{value_type.__name__}"
+
+
+def serialize_workflow_event(event: WorkflowEvent[Any]) -> dict[str, Any]:
+    """Serialize a :class:`WorkflowEvent` to a JSON-compatible dict.
+
+    Carries a workflow event from the durable activity, through the orchestration
+    custom status, to a streaming client. The data payload is encoded with
+    :func:`serialize_value` so typed objects survive the round trip;
+    :func:`deserialize_workflow_event` reverses it into a ``WorkflowEvent`` so
+    callers never handle checkpoint-marker dicts directly.
+
+    Args:
+        event: The workflow event to serialize.
+
+    Returns:
+        A JSON-serializable dict with the event ``type`` and the fields needed to
+        reconstruct it.
+    """
+    serialized: dict[str, Any] = {"type": event.type}
+    if event.executor_id is not None:
+        serialized["executor_id"] = event.executor_id
+    if event.data is not None:
+        serialized["data"] = serialize_value(event.data)
+    if event.type == "request_info":
+        # request_type is omitted: deserialize_workflow_event rebuilds the event via
+        # WorkflowEvent.request_info, which derives it from the data payload.
+        serialized["request_id"] = event.request_id
+        serialized["source_executor_id"] = event.source_executor_id
+        serialized["response_type"] = _type_key(event.response_type)
+    return serialized
+
+
+def deserialize_workflow_event(serialized: dict[str, Any]) -> WorkflowEvent[Any]:
+    """Reconstruct a :class:`WorkflowEvent` from :func:`serialize_workflow_event` output.
+
+    ``serialized`` must originate from the workflow's own orchestration custom
+    status (trusted durable storage); its encoded payload is decoded with
+    :func:`deserialize_value`. Never pass untrusted external input here.
+
+    Args:
+        serialized: A dict previously produced by :func:`serialize_workflow_event`,
+            optionally augmented with an ``iteration`` key by the orchestrator.
+
+    Returns:
+        The reconstructed workflow event with its data payload restored.
+    """
+    event_type = cast(WorkflowEventType, serialized["type"])
+    payload = deserialize_value(serialized["data"]) if "data" in serialized else None
+
+    if event_type == "request_info":
+        response_key = serialized.get("response_type")
+        response_type = resolve_type(response_key) if response_key else None
+        event: WorkflowEvent[Any] = WorkflowEvent.request_info(
+            request_id=cast(str, serialized["request_id"]),
+            source_executor_id=cast(str, serialized["source_executor_id"]),
+            request_data=payload,
+            response_type=response_type or object,
+        )
+    else:
+        event = WorkflowEvent(event_type, data=payload, executor_id=serialized.get("executor_id"))
+
+    iteration = serialized.get("iteration")
+    if iteration is not None:
+        event.iteration = iteration
+    return event
+
+
+# ============================================================================
+# HITL Type Reconstruction
+# ============================================================================
+
+
+def reconstruct_to_type(value: Any, target_type: type) -> Any:
+    """Reconstruct a value to a known target type.
+
+    Used for HITL responses where external data (without checkpoint type markers)
+    needs to be reconstructed to a specific type determined by the response_type hint.
+
+    Tries strategies in order:
+    1. Return as-is if already the correct type
+    2. deserialize_value (for data with any type markers)
+    3. Pydantic model_validate (for Pydantic models)
+    4. Dataclass constructor (for dataclasses)
+
+    Args:
+        value: The value to reconstruct (typically a dict from JSON)
+        target_type: The expected type to reconstruct to
+
+    Returns:
+        Reconstructed value if possible, otherwise the original value
+    """
+    if value is None:
+        return None
+
+    with suppress(TypeError):
+        if isinstance(value, target_type):
+            return value
+
+    if not isinstance(value, dict):
+        return value
+
+    # Try decoding if data has pickle markers (from checkpoint encoding).
+    # NOTE: This function is general-purpose.  Callers that handle untrusted
+    # data (e.g. HITL responses) MUST call strip_pickle_markers() before
+    # passing data here.  See _deserialize_hitl_response in orchestrator.py.
+    decoded = deserialize_value(value)
+    if not isinstance(decoded, dict):
+        return decoded
+
+    # Try Pydantic model validation (for unmarked dicts, e.g., external HITL data)
+    if issubclass(target_type, BaseModel):
+        try:
+            return target_type.model_validate(value)
+        except Exception:
+            logger.debug("Could not validate Pydantic model %s", target_type)
+            return value  # type: ignore[return-value]
+
+    # Try dataclass construction (for unmarked dicts, e.g., external HITL data)
+    if is_dataclass(target_type) and isinstance(target_type, type):  # type: ignore
+        try:
+            return target_type(**value)
+        except Exception:
+            logger.debug("Could not construct dataclass %s", target_type)
+
+    return value  # type: ignore[return-value]

python/packages/durabletask/pyproject.toml

@@ -0,0 +1,110 @@
+[project]
+name = "agent-framework-durabletask"
+description = "Durable Task integration for Microsoft Agent Framework."
+authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
+readme = "README.md"
+requires-python = ">=3.10"
+version = "1.0.0b260709"
+license-files = ["LICENSE"]
+urls.homepage = "https://aka.ms/agent-framework"
+urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
+urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
+urls.issues = "https://github.com/microsoft/agent-framework/issues"
+classifiers = [
+  "License :: OSI Approved :: MIT License",
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Typing :: Typed",
+]
+dependencies = [
+    "agent-framework-core>=1.11.0,<2",
+    "durabletask>=1.5.0,<2",
+    "durabletask-azuremanaged>=1.4.0,<2",
+    "python-dateutil>=2.8.0,<3",
+]
+
+[dependency-groups]
+dev = [
+    "types-python-dateutil==2.9.0.20260518",
+]
+
+[tool.uv]
+prerelease = "if-necessary-or-explicit"
+environments = [
+    "sys_platform == 'darwin'",
+    "sys_platform == 'linux'",
+    "sys_platform == 'win32'"
+]
+
+[tool.uv-dynamic-versioning]
+fallback-version = "0.0.0"
+
+[tool.pytest.ini_options]
+testpaths = 'tests'
+pythonpath = ["tests/integration_tests"]
+addopts = "-ra -q -r fEX"
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+filterwarnings = [
+    "ignore:Support for class-based `config` is deprecated:DeprecationWarning:pydantic.*"
+]
+timeout = 120
+markers = [
+    "integration: marks tests as integration tests",
+    "integration_test: marks tests as integration tests (alternative marker)",
+    "sample: marks tests as sample tests",
+    "requires_azure_openai: marks tests that require Azure OpenAI",
+    "requires_dts: marks tests that require Durable Task Scheduler",
+    "requires_redis: marks tests that require Redis"
+]
+
+[tool.ruff]
+extend = "../../pyproject.toml"
+
+[tool.coverage.run]
+omit = [
+    "**/__init__.py"
+]
+
+[tool.pyright]
+extends = "../../pyproject.toml"
+include = ["agent_framework_durabletask"]
+
+[tool.mypy]
+plugins = ['pydantic.mypy']
+strict = true
+python_version = "3.10"
+ignore_missing_imports = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+check_untyped_defs = true
+warn_return_any = true
+show_error_codes = true
+warn_unused_ignores = false
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+
+[tool.bandit]
+targets = ["agent_framework_durabletask"]
+exclude_dirs = ["tests"]
+
+[tool.poe]
+executor.type = "uv"
+include = "../../shared_tasks.toml"
+
+[tool.poe.tasks.mypy]
+help = "Run MyPy for this package."
+cmd = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_durabletask"
+
+[tool.poe.tasks.test]
+help = "Run the default unit test suite for this package."
+cmd = 'pytest -m "not integration" --cov=agent_framework_durabletask --cov-report=term-missing:skip-covered tests'
+
+[build-system]
+requires = ["flit-core >= 3.11,<4.0"]
+build-backend = "flit_core.buildapi"

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