Python: Complete durableagent package (#3058)

* Add worker and clients

* Clean code and refactor common code

* Implement sample

* Add sample

* Update readmes

* Fix tests

* Fix tests

* Update requirements

* Fix typo

* Address comments

* use response.text

python/packages/durabletask/agent_framework_durabletask/__init__.py

@@ -3,6 +3,7 @@
 """Durable Task integration for Microsoft Agent Framework."""
 
 from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
+from ._client import DurableAIAgentClient
 from ._constants import (
     DEFAULT_MAX_POLL_RETRIES,
     DEFAULT_POLL_INTERVAL_SECONDS,
@@ -41,7 +42,12 @@ from ._durable_agent_state import (
     DurableAgentStateUsageContent,
 )
 from ._entities import AgentEntity, AgentEntityStateProviderMixin
-from ._models import RunRequest, serialize_response_format
+from ._executors import DurableAgentExecutor
+from ._models import AgentSessionId, DurableAgentThread, RunRequest
+from ._orchestration_context import DurableAIAgentOrchestrationContext
+from ._response_utils import ensure_response_format, load_agent_response
+from ._shim import DurableAIAgent
+from ._worker import DurableAIAgentWorker
 
 __all__ = [
     "DEFAULT_MAX_POLL_RETRIES",
@@ -58,8 +64,14 @@ __all__ = [
     "AgentEntity",
     "AgentEntityStateProviderMixin",
     "AgentResponseCallbackProtocol",
+    "AgentSessionId",
     "ApiResponseFields",
     "ContentTypes",
+    "DurableAIAgent",
+    "DurableAIAgentClient",
+    "DurableAIAgentOrchestrationContext",
+    "DurableAIAgentWorker",
+    "DurableAgentExecutor",
     "DurableAgentState",
     "DurableAgentStateContent",
     "DurableAgentStateData",
@@ -80,7 +92,10 @@ __all__ = [
     "DurableAgentStateUriContent",
     "DurableAgentStateUsage",
     "DurableAgentStateUsageContent",
+    "DurableAgentThread",
+    "DurableAgentThread",
     "DurableStateFields",
     "RunRequest",
-    "serialize_response_format",
+    "ensure_response_format",
+    "load_agent_response",
 ]

python/packages/durabletask/agent_framework_durabletask/_client.py

@@ -0,0 +1,90 @@
+# 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
+
+from agent_framework import AgentRunResponse, get_logger
+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 = get_logger("agent_framework.durabletask.client")
+
+
+class DurableAIAgentClient(DurableAgentProvider[AgentRunResponse]):
+    """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_durabletask 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[AgentRunResponse]:
+        """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/_durable_agent_state.py

@@ -53,7 +53,7 @@ from agent_framework import (
 )
 from dateutil import parser as date_parser
 
-from ._constants import ApiResponseFields, ContentTypes, DurableStateFields
+from ._constants import ContentTypes, DurableStateFields
 from ._models import RunRequest, serialize_response_format
 
 logger = get_logger("agent_framework.durabletask.durable_agent_state")
@@ -452,7 +452,7 @@ class DurableAgentState:
         """Get the count of conversation entries (requests + responses)."""
         return len(self.data.conversation_history)
 
-    def try_get_agent_response(self, correlation_id: str) -> dict[str, Any] | None:
+    def try_get_agent_response(self, correlation_id: str) -> AgentRunResponse | None:
         """Try to get an agent response by correlation ID.
 
         This method searches the conversation history for a response entry matching the given
@@ -474,14 +474,8 @@ class DurableAgentState:
         for entry in self.data.conversation_history:
             if entry.correlation_id == correlation_id and isinstance(entry, DurableAgentStateResponse):
                 # Found the entry, extract response data
-                # Get the text content from assistant messages only
-                content = "\n".join(message.text for message in entry.messages if message.text)
+                return DurableAgentStateResponse.to_run_response(entry)
 
-                return {
-                    ApiResponseFields.CONTENT: content,
-                    ApiResponseFields.MESSAGE_COUNT: self.message_count,
-                    ApiResponseFields.CORRELATION_ID: correlation_id,
-                }
         return None
 
 
@@ -705,6 +699,21 @@ class DurableAgentStateResponse(DurableAgentStateEntry):
             usage=DurableAgentStateUsage.from_usage(response.usage_details),
         )
 
+    @staticmethod
+    def to_run_response(
+        response_entry: DurableAgentStateResponse,
+    ) -> AgentRunResponse:
+        """Converts a DurableAgentStateResponse back to an AgentRunResponse."""
+        messages = [m.to_chat_message() for m in response_entry.messages]
+
+        usage_details = response_entry.usage.to_usage_details() if response_entry.usage is not None else UsageDetails()
+
+        return AgentRunResponse(
+            created_at=response_entry.created_at.isoformat(),
+            messages=messages,
+            usage_details=usage_details,
+        )
+
 
 class DurableAgentStateMessage:
     """Represents a message within a conversation history entry.
@@ -1214,14 +1223,24 @@ class DurableAgentStateUsage:
             input_token_count=usage.input_token_count,
             output_token_count=usage.output_token_count,
             total_token_count=usage.total_token_count,
+            extensionData=usage.additional_counts,
         )
 
     def to_usage_details(self) -> UsageDetails:
         # Convert back to AI SDK UsageDetails
+        extension_data: dict[str, int] = {}
+        if self.extensionData is not None:
+            for k, v in self.extensionData.items():
+                try:
+                    extension_data[k] = int(v)
+                except (ValueError, TypeError):
+                    continue
+
         return UsageDetails(
             input_token_count=self.input_token_count,
             output_token_count=self.output_token_count,
             total_token_count=self.total_token_count,
+            **extension_data,
         )
 
 

python/packages/durabletask/agent_framework_durabletask/_entities.py

@@ -128,7 +128,7 @@ class AgentEntity:
     ) -> AgentRunResponse:
         """Execute the agent with a message."""
         if isinstance(request, str):
-            run_request = RunRequest(message=request, role=Role.USER)
+            run_request = RunRequest.from_json(request)
         elif isinstance(request, dict):
             run_request = RunRequest.from_dict(request)
         else:
@@ -139,8 +139,6 @@ class AgentEntity:
         correlation_id = run_request.correlation_id
         if not thread_id:
             raise ValueError("Entity State Provider must provide a thread_id")
-        if not correlation_id:
-            raise ValueError("RunRequest must include a correlation_id")
         response_format = run_request.response_format
         enable_tool_calls = run_request.enable_tool_calls
 

python/packages/durabletask/agent_framework_durabletask/_executors.py

@@ -0,0 +1,460 @@
+# 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 time
+import uuid
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+from typing import Any, Generic, TypeVar
+
+from agent_framework import AgentRunResponse, AgentThread, ChatMessage, ErrorContent, Role, get_logger
+from durabletask.client import TaskHubGrpcClient
+from durabletask.entities import EntityInstanceId
+from durabletask.task import 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, DurableAgentThread, RunRequest
+from ._response_utils import ensure_response_format, load_agent_response
+
+logger = get_logger("agent_framework.durabletask.executors")
+
+# TypeVar for the task type returned by executors
+TaskT = TypeVar("TaskT")
+
+
+class DurableAgentTask(CompositeTask[AgentRunResponse]):
+    """A custom Task that wraps entity calls and provides typed AgentRunResponse results.
+
+    This task wraps the underlying entity call task and intercepts its completion
+    to convert the raw result into a typed AgentRunResponse object.
+    """
+
+    def __init__(
+        self,
+        entity_task: Task[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])  # type: ignore[misc]
+
+    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
+            self._exception = task.get_exception()
+            self._is_complete = True
+            if self._parent is not None:
+                self._parent.on_child_completed(self)
+            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 AgentRunResponse as this task's result
+            self._result = response
+            self._is_complete = True
+
+            if self._parent is not None:
+                self._parent.on_child_completed(self)
+
+        except Exception:
+            logger.exception(
+                "[DurableAgentTask] Failed to convert result for correlation_id: %s",
+                self._correlation_id,
+            )
+            raise
+
+
+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,
+        thread: AgentThread | None = None,
+    ) -> TaskT:
+        """Execute the durable agent.
+
+        Returns:
+            TaskT: The task type specific to this executor implementation
+        """
+        raise NotImplementedError
+
+    def get_new_thread(self, agent_name: str, **kwargs: Any) -> DurableAgentThread:
+        """Create a new DurableAgentThread with random session ID."""
+        session_id = self._create_session_id(agent_name)
+        return DurableAgentThread.from_session_id(session_id, **kwargs)
+
+    def _create_session_id(
+        self,
+        agent_name: str,
+        thread: AgentThread | None = None,
+    ) -> AgentSessionId:
+        """Create the AgentSessionId for the execution."""
+        if isinstance(thread, DurableAgentThread) and thread.session_id is not None:
+            return thread.session_id
+        # Create new session ID - either no thread provided or it's a regular AgentThread
+        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,
+        response_format: type[BaseModel] | None,
+        enable_tool_calls: bool,
+    ) -> RunRequest:
+        """Create a RunRequest for the given parameters."""
+        correlation_id = self.generate_unique_id()
+        return RunRequest(
+            message=message,
+            response_format=response_format,
+            enable_tool_calls=enable_tool_calls,
+            correlation_id=correlation_id,
+        )
+
+
+class ClientAgentExecutor(DurableAgentExecutor[AgentRunResponse]):
+    """Execution strategy for external clients.
+
+    Note: Returns AgentRunResponse 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,
+        thread: AgentThread | None = None,
+    ) -> AgentRunResponse:
+        """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
+            thread: Optional conversation thread (creates new if not provided)
+
+        Returns:
+            AgentRunResponse: The agent's response after execution completes
+        """
+        # Signal the entity with the request
+        entity_id = self._signal_agent_entity(agent_name, run_request, thread)
+
+        # 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,
+        thread: AgentThread | 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
+            thread: Optional conversation thread
+
+        Returns:
+            entity_id
+        """
+        # Get or create session ID
+        session_id = self._create_session_id(agent_name, thread)
+
+        # 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,
+    ) -> AgentRunResponse | 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):
+            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: AgentRunResponse | None,
+        response_format: type[BaseModel] | None,
+        correlation_id: str,
+    ) -> AgentRunResponse:
+        """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:
+            AgentRunResponse 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 = ChatMessage(
+                    role=Role.SYSTEM,
+                    contents=[
+                        ErrorContent(
+                            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 = ChatMessage(
+                role=Role.SYSTEM,
+                contents=[
+                    ErrorContent(
+                        message=f"Timeout waiting for agent response after {self.max_poll_retries} attempts",
+                        error_code="response_timeout",
+                    )
+                ],
+            )
+
+        return AgentRunResponse(
+            messages=[error_message],
+            created_at=datetime.now(timezone.utc).isoformat(),
+        )
+
+    def _poll_entity_for_response(
+        self,
+        entity_id: EntityInstanceId,
+        correlation_id: str,
+    ) -> AgentRunResponse | 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 AgentRunResponse, 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 get_run_request(
+        self,
+        message: str,
+        response_format: type[BaseModel] | None,
+        enable_tool_calls: bool,
+    ) -> RunRequest:
+        """Get the current run request from the orchestration context.
+
+        Returns:
+            RunRequest: The current run request
+        """
+        request = super().get_run_request(
+            message,
+            response_format,
+            enable_tool_calls,
+        )
+        request.orchestration_id = self._context.instance_id
+        return request
+
+    def run_durable_agent(
+        self,
+        agent_name: str,
+        run_request: RunRequest,
+        thread: AgentThread | 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
+            thread: Optional conversation thread (creates new if not provided)
+
+        Returns:
+            DurableAgentTask: A task wrapping the entity call that yields AgentRunResponse
+        """
+        # Resolve session
+        session_id = self._create_session_id(agent_name, thread)
+
+        # 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,
+        )
+
+        # Call the entity and get the underlying task
+        entity_task: Task[Any] = self._context.call_entity(entity_id, "run", run_request.to_dict())  # type: ignore
+
+        # 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

@@ -8,12 +8,15 @@ This module defines the request and response models used by the framework.
 from __future__ import annotations
 
 import inspect
+import json
+import uuid
+from collections.abc import MutableMapping
 from dataclasses import dataclass
-from datetime import datetime
+from datetime import datetime, timezone
 from importlib import import_module
 from typing import TYPE_CHECKING, Any, cast
 
-from agent_framework import Role
+from agent_framework import AgentThread, Role
 
 from ._constants import REQUEST_RESPONSE_FORMAT_TEXT
 
@@ -101,38 +104,38 @@ class RunRequest:
         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
-        correlation_id: Optional correlation ID for tracking the response to this specific request
+        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
     """
 
     message: str
     request_response_format: str
+    correlation_id: str
     role: Role = Role.USER
     response_format: type[BaseModel] | None = None
     enable_tool_calls: bool = True
-    correlation_id: str | None = None
     created_at: datetime | None = None
     orchestration_id: str | None = None
 
     def __init__(
         self,
         message: str,
+        correlation_id: str,
         request_response_format: str = REQUEST_RESPONSE_FORMAT_TEXT,
         role: Role | str | None = Role.USER,
         response_format: type[BaseModel] | None = None,
         enable_tool_calls: bool = True,
-        correlation_id: str | None = None,
         created_at: datetime | None = None,
         orchestration_id: str | 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.correlation_id = correlation_id
-        self.created_at = created_at
+        self.created_at = created_at if created_at is not None else datetime.now(tz=timezone.utc)
         self.orchestration_id = orchestration_id
 
     @staticmethod
@@ -154,11 +157,10 @@ class RunRequest:
             "enable_tool_calls": self.enable_tool_calls,
             "role": self.role.value,
             "request_response_format": self.request_response_format,
+            "correlationId": self.correlation_id,
         }
         if self.response_format:
             result["response_format"] = serialize_response_format(self.response_format)
-        if self.correlation_id:
-            result["correlationId"] = self.correlation_id
         if self.created_at:
             result["created_at"] = self.created_at.isoformat()
         if self.orchestration_id:
@@ -166,6 +168,16 @@ class RunRequest:
 
         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."""
@@ -176,13 +188,120 @@ class RunRequest:
             except ValueError:
                 created_at = None
 
+        correlation_id = data.get("correlationId")
+        if not correlation_id:
+            raise ValueError("correlationId is required in RunRequest data")
+
         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")),
             enable_tool_calls=data.get("enable_tool_calls", True),
-            correlation_id=data.get("correlationId"),
             created_at=created_at,
             orchestration_id=data.get("orchestrationId"),
         )
+
+
+@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) -> AgentSessionId:
+        if not session_id_string.startswith("@"):
+            raise ValueError(f"Invalid agent session ID format: {session_id_string}")
+
+        parts = session_id_string[1:].split("@", 1)
+        if len(parts) != 2:
+            raise ValueError(f"Invalid agent session ID format: {session_id_string}")
+
+        return AgentSessionId(name=parts[0], key=parts[1])
+
+
+class DurableAgentThread(AgentThread):
+    """Durable agent thread that tracks the owning :class:`AgentSessionId`."""
+
+    _SERIALIZED_SESSION_ID_KEY = "durable_session_id"
+
+    def __init__(
+        self,
+        *,
+        session_id: AgentSessionId | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._session_id: AgentSessionId | None = session_id
+
+    @property
+    def session_id(self) -> AgentSessionId | None:
+        return self._session_id
+
+    @session_id.setter
+    def session_id(self, value: AgentSessionId | None) -> None:
+        self._session_id = value
+
+    @classmethod
+    def from_session_id(
+        cls,
+        session_id: AgentSessionId,
+        **kwargs: Any,
+    ) -> DurableAgentThread:
+        return cls(session_id=session_id, **kwargs)
+
+    async def serialize(self, **kwargs: Any) -> dict[str, Any]:
+        state = await super().serialize(**kwargs)
+        if self._session_id is not None:
+            state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id)
+        return state
+
+    @classmethod
+    async def deserialize(
+        cls,
+        serialized_thread_state: MutableMapping[str, Any],
+        *,
+        message_store: Any = None,
+        **kwargs: Any,
+    ) -> DurableAgentThread:
+        state_payload = dict(serialized_thread_state)
+        session_id_value = state_payload.pop(cls._SERIALIZED_SESSION_ID_KEY, None)
+        thread = await super().deserialize(
+            state_payload,
+            message_store=message_store,
+            **kwargs,
+        )
+        if not isinstance(thread, DurableAgentThread):
+            raise TypeError("Deserialized thread is not a DurableAgentThread instance")
+
+        if session_id_value is None:
+            return thread
+
+        if not isinstance(session_id_value, str):
+            raise ValueError("durable_session_id must be a string when present in serialized state")
+
+        thread.session_id = AgentSessionId.parse(session_id_value)
+        return thread

python/packages/durabletask/agent_framework_durabletask/_orchestration_context.py

@@ -0,0 +1,75 @@
+# 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
+
+from agent_framework import get_logger
+from durabletask.task import OrchestrationContext
+
+from ._executors import DurableAgentTask, OrchestrationAgentExecutor
+from ._shim import DurableAgentProvider, DurableAIAgent
+
+logger = get_logger("agent_framework.durabletask.orchestration_context")
+
+
+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_durabletask 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,62 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Shared utilities for handling AgentRunResponse parsing and validation."""
+
+from typing import Any
+
+from agent_framework import AgentRunResponse, get_logger
+from pydantic import BaseModel
+
+logger = get_logger("agent_framework.durabletask.response_utils")
+
+
+def load_agent_response(agent_response: AgentRunResponse | dict[str, Any] | None) -> AgentRunResponse:
+    """Convert raw payloads into AgentRunResponse instance.
+
+    Args:
+        agent_response: The response to convert, can be an AgentRunResponse, dict, or None
+
+    Returns:
+        AgentRunResponse: 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, AgentRunResponse):
+        return agent_response
+    if isinstance(agent_response, dict):
+        logger.debug("[load_agent_response] Converting dict payload using AgentRunResponse.from_dict")
+        return AgentRunResponse.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: AgentRunResponse,
+) -> None:
+    """Ensure the AgentRunResponse 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 AgentRunResponse object to validate and parse
+    """
+    if response_format is not None and not isinstance(response.value, response_format):
+        response.try_parse_value(response_format)
+
+        logger.debug(
+            "[ensure_response_format] Loaded AgentRunResponse.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,185 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Durable Agent Shim for Durable Task Framework.
+
+This module provides the DurableAIAgent shim that implements AgentProtocol
+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 collections.abc import AsyncIterator
+from typing import Any, Generic, TypeVar
+
+from agent_framework import AgentProtocol, AgentRunResponseUpdate, AgentThread, ChatMessage
+from pydantic import BaseModel
+
+from ._executors import DurableAgentExecutor
+from ._models import DurableAgentThread
+
+# 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(AgentProtocol, Generic[TaskT]):
+    """A durable agent proxy that delegates execution to the provider.
+
+    This class implements AgentProtocol but with one critical difference:
+    - AgentProtocol.run() returns a Coroutine (async, must await)
+    - DurableAIAgent.run() returns TaskT (sync Task object - must yield
+        or the AgentRunResponse 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., AgentRunResponse, DurableAgentTask, AgentTask)
+    """
+
+    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
+        self._id = agent_id if agent_id is not None else name
+        self._display_name = name
+        self._description = f"Durable agent proxy for {name}"
+
+    @property
+    def id(self) -> str:
+        """Get the unique identifier for this agent."""
+        return self._id
+
+    @property
+    def name(self) -> str | None:
+        """Get the name of the agent."""
+        return self._name
+
+    @property
+    def display_name(self) -> str:
+        """Get the display name of the agent."""
+        return self._display_name
+
+    @property
+    def description(self) -> str | None:
+        """Get the description of the agent."""
+        return self._description
+
+    def run(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        response_format: type[BaseModel] | None = None,
+        enable_tool_calls: bool = True,
+    ) -> TaskT:
+        """Execute the agent via the injected provider.
+
+        Note:
+            This method overrides AgentProtocol.run() with a different return type:
+            - AgentProtocol.run() returns Coroutine[Any, Any, AgentRunResponse] (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
+        """
+        message_str = self._normalize_messages(messages)
+
+        run_request = self._executor.get_run_request(
+            message=message_str,
+            response_format=response_format,
+            enable_tool_calls=enable_tool_calls,
+        )
+
+        return self._executor.run_durable_agent(
+            agent_name=self._name,
+            run_request=run_request,
+            thread=thread,
+        )
+
+    def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[AgentRunResponseUpdate]:
+        """Run the agent with streaming (not supported for durable agents).
+
+        Args:
+            messages: The message(s) to send to the agent
+            thread: Optional agent thread for conversation context
+            **kwargs: Additional arguments
+
+        Raises:
+            NotImplementedError: Streaming is not supported for durable agents
+        """
+        raise NotImplementedError("Streaming is not supported for durable agents")
+
+    def get_new_thread(self, **kwargs: Any) -> DurableAgentThread:
+        """Create a new agent thread via the provider."""
+        return self._executor.get_new_thread(self._name, **kwargs)
+
+    def _normalize_messages(self, messages: str | ChatMessage | list[str] | list[ChatMessage] | None) -> str:
+        """Convert supported message inputs to a single string.
+
+        Args:
+            messages: The messages to normalize
+
+        Returns:
+            A single string representation of the messages
+        """
+        if messages is None:
+            return ""
+        if isinstance(messages, str):
+            return messages
+        if isinstance(messages, ChatMessage):
+            return messages.text or ""
+        if isinstance(messages, list):
+            if not messages:
+                return ""
+            first_item = messages[0]
+            if isinstance(first_item, str):
+                return "\n".join(messages)  # type: ignore[arg-type]
+            # List of ChatMessage
+            return "\n".join([msg.text or "" for msg in messages])  # type: ignore[union-attr]
+        return ""

python/packages/durabletask/agent_framework_durabletask/_worker.py

@@ -0,0 +1,218 @@
+# 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.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from agent_framework import AgentProtocol, get_logger
+from durabletask.worker import TaskHubGrpcWorker
+
+from ._callbacks import AgentResponseCallbackProtocol
+from ._entities import AgentEntity, DurableTaskEntityStateProvider
+
+logger = get_logger("agent_framework.durabletask.worker")
+
+
+class DurableAIAgentWorker:
+    """Wrapper for durabletask worker that enables agent registration.
+
+    This class wraps an existing TaskHubGrpcWorker instance and provides
+    a convenient interface for registering agents as durable entities.
+
+    Example:
+        ```python
+        from durabletask import TaskHubGrpcWorker
+        from agent_framework import ChatAgent
+        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
+        my_agent = ChatAgent(chat_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, AgentProtocol] = {}
+        logger.debug("[DurableAIAgentWorker] Initialized with worker type: %s", type(worker).__name__)
+
+    def add_agent(
+        self,
+        agent: AgentProtocol,
+        callback: AgentResponseCallbackProtocol | 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-{agent_name}".
+
+        Args:
+            agent: The agent to register (must have a name)
+            callback: Optional callback for this specific agent (overrides worker-level callback)
+
+        Raises:
+            ValueError: If the agent doesn't have a name or is already registered
+        """
+        agent_name = agent.name
+        if not agent_name:
+            raise ValueError("Agent must have a name to be registered")
+
+        if agent_name in self._registered_agents:
+            raise ValueError(f"Agent '{agent_name}' is already registered")
+
+        logger.info("[DurableAIAgentWorker] Registering agent: %s as entity: dafx-%s", agent_name, agent_name)
+
+        # Store the agent reference
+        self._registered_agents[agent_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)
+
+        # Register the entity class with the worker
+        # The worker.add_entity method takes a class
+        entity_registered: str = self._worker.add_entity(entity_class)  # pyright: ignore[reportUnknownMemberType]
+
+        logger.debug(
+            "[DurableAIAgentWorker] Successfully registered entity class %s for agent: %s",
+            entity_registered,
+            agent_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())
+
+    def __create_agent_entity(
+        self,
+        agent: AgentProtocol,
+        callback: AgentResponseCallbackProtocol | 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
+
+        Returns:
+            A new DurableEntity subclass configured for this agent
+        """
+        agent_name = 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:
+                    AgentRunResponse as dict
+                """
+                logger.debug("[ConfiguredAgentEntity.run] Executing agent: %s", agent_name)
+                # Get or create event loop for async execution
+                try:
+                    loop = asyncio.get_event_loop()
+                except RuntimeError:
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+
+                # Run the async agent execution synchronously
+                if loop.is_running():
+                    # If loop is already running (shouldn't happen in entity context),
+                    # create a temporary loop
+                    temp_loop = asyncio.new_event_loop()
+                    try:
+                        response = temp_loop.run_until_complete(self._agent_entity.run(request))
+                    finally:
+                        temp_loop.close()
+                else:
+                    response = loop.run_until_complete(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