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/tests/test_agent_session_id.py

@@ -0,0 +1,272 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for AgentSessionId and DurableAgentThread."""
+
+import pytest
+from agent_framework import AgentThread
+
+from agent_framework_durabletask._models import AgentSessionId, DurableAgentThread
+
+
+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"
+
+
+class TestDurableAgentThread:
+    """Test suite for DurableAgentThread."""
+
+    def test_init_with_session_id(self) -> None:
+        """Test DurableAgentThread initialization with session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        thread = DurableAgentThread(session_id=session_id)
+
+        assert thread.session_id is not None
+        assert thread.session_id == session_id
+
+    def test_init_without_session_id(self) -> None:
+        """Test DurableAgentThread initialization without session ID."""
+        thread = DurableAgentThread()
+
+        assert thread.session_id is None
+
+    def test_session_id_setter(self) -> None:
+        """Test setting a session ID to an existing thread."""
+        thread = DurableAgentThread()
+        assert thread.session_id is None
+
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        thread.session_id = session_id
+
+        assert thread.session_id is not None
+        assert thread.session_id == session_id
+        assert thread.session_id.name == "TestAgent"
+
+    def test_from_session_id(self) -> None:
+        """Test creating DurableAgentThread from session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        thread = DurableAgentThread.from_session_id(session_id)
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.session_id is not None
+        assert thread.session_id == session_id
+        assert thread.session_id.name == "TestAgent"
+        assert thread.session_id.key == "test-key"
+
+    def test_from_session_id_with_service_thread_id(self) -> None:
+        """Test creating DurableAgentThread with service thread ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        thread = DurableAgentThread.from_session_id(session_id, service_thread_id="service-123")
+
+        assert thread.session_id is not None
+        assert thread.session_id == session_id
+        assert thread.service_thread_id == "service-123"
+
+    async def test_serialize_with_session_id(self) -> None:
+        """Test serialization includes session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key")
+        thread = DurableAgentThread(session_id=session_id)
+
+        serialized = await thread.serialize()
+
+        assert isinstance(serialized, dict)
+        assert "durable_session_id" in serialized
+        assert serialized["durable_session_id"] == "@TestAgent@test-key"
+
+    async def test_serialize_without_session_id(self) -> None:
+        """Test serialization without session ID."""
+        thread = DurableAgentThread()
+
+        serialized = await thread.serialize()
+
+        assert isinstance(serialized, dict)
+        assert "durable_session_id" not in serialized
+
+    async def test_deserialize_with_session_id(self) -> None:
+        """Test deserialization restores session ID."""
+        serialized = {
+            "service_thread_id": "thread-123",
+            "durable_session_id": "@TestAgent@test-key",
+        }
+
+        thread = await DurableAgentThread.deserialize(serialized)
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.session_id is not None
+        assert thread.session_id.name == "TestAgent"
+        assert thread.session_id.key == "test-key"
+        assert thread.service_thread_id == "thread-123"
+
+    async def test_deserialize_without_session_id(self) -> None:
+        """Test deserialization without session ID."""
+        serialized = {
+            "service_thread_id": "thread-456",
+        }
+
+        thread = await DurableAgentThread.deserialize(serialized)
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.session_id is None
+        assert thread.service_thread_id == "thread-456"
+
+    async def test_round_trip_serialization(self) -> None:
+        """Test round-trip serialization preserves session ID."""
+        session_id = AgentSessionId(name="TestAgent", key="test-key-789")
+        original = DurableAgentThread(session_id=session_id)
+
+        serialized = await original.serialize()
+        restored = await DurableAgentThread.deserialize(serialized)
+
+        assert isinstance(restored, DurableAgentThread)
+        assert restored.session_id is not None
+        assert restored.session_id.name == session_id.name
+        assert restored.session_id.key == session_id.key
+
+    async def test_deserialize_invalid_session_id_type(self) -> None:
+        """Test deserialization with invalid session ID type raises error."""
+        serialized = {
+            "service_thread_id": "thread-123",
+            "durable_session_id": 12345,  # Invalid type
+        }
+
+        with pytest.raises(ValueError, match="durable_session_id must be a string"):
+            await DurableAgentThread.deserialize(serialized)
+
+
+class TestAgentThreadCompatibility:
+    """Test suite for compatibility between AgentThread and DurableAgentThread."""
+
+    async def test_agent_thread_serialize(self) -> None:
+        """Test that base AgentThread can be serialized."""
+        thread = AgentThread()
+
+        serialized = await thread.serialize()
+
+        assert isinstance(serialized, dict)
+        assert "service_thread_id" in serialized
+
+    async def test_agent_thread_deserialize(self) -> None:
+        """Test that base AgentThread can be deserialized."""
+        thread = AgentThread()
+        serialized = await thread.serialize()
+
+        restored = await AgentThread.deserialize(serialized)
+
+        assert isinstance(restored, AgentThread)
+        assert restored.service_thread_id == thread.service_thread_id
+
+    async def test_durable_thread_is_agent_thread(self) -> None:
+        """Test that DurableAgentThread is an AgentThread."""
+        thread = DurableAgentThread()
+
+        assert isinstance(thread, AgentThread)
+        assert isinstance(thread, DurableAgentThread)
+
+
+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@")
+
+    async def test_thread_with_session_preserves_on_serialization(self) -> None:
+        """Test that thread with session ID preserves it through serialization."""
+        session_id = AgentSessionId(name="TestAgent", key="preserved-key")
+        thread = DurableAgentThread.from_session_id(session_id)
+
+        # Serialize and deserialize
+        serialized = await thread.serialize()
+        restored = await DurableAgentThread.deserialize(serialized)
+
+        # Session ID should be preserved
+        assert restored.session_id is not None
+        assert restored.session_id.name == "TestAgent"
+        assert restored.session_id.key == "preserved-key"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_client.py

@@ -0,0 +1,142 @@
+# 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 AgentProtocol
+
+from agent_framework_durabletask import DurableAgentThread, 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, AgentProtocol)
+
+    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_threads(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify agent from client can create DurableAgentThread instances."""
+        agent = agent_client.get_agent("assistant")
+
+        thread = agent.get_new_thread()
+
+        assert isinstance(thread, DurableAgentThread)
+
+    def test_client_agent_thread_with_parameters(self, agent_client: DurableAIAgentClient) -> None:
+        """Verify agent can create threads with custom parameters."""
+        agent = agent_client.get_agent("assistant")
+
+        thread = agent.get_new_thread(service_thread_id="client-session-123")
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.service_thread_id == "client-session-123"
+
+
+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

@@ -112,20 +112,21 @@ 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 preserves None created_at instead of defaulting to current time.
+        """Test from_run_request handles auto-populated created_at from RunRequest.
 
-        When a RunRequest has no created_at value, the resulting DurableAgentStateMessage
-        should also have None for created_at, not default to current UTC time.
+        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,  # Explicitly None
+            created_at=None,  # RunRequest will default this to current time
         )
 
         durable_message = DurableAgentStateMessage.from_run_request(run_request)
 
-        assert durable_message.created_at is None
+        # 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."""

python/packages/durabletask/tests/test_executors.py

@@ -0,0 +1,320 @@
+# 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 AgentRunResponse, Role
+from durabletask.entities import EntityInstanceId
+from durabletask.task import Task
+from pydantic import BaseModel
+
+from agent_framework_durabletask import DurableAgentThread
+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."""
+    return Mock(spec=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)
+    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",
+    }
+
+
+class TestExecutorThreadCreation:
+    """Test that executors properly create DurableAgentThread with parameters."""
+
+    def test_client_executor_creates_durable_thread(self, mock_client: Mock) -> None:
+        """Verify ClientAgentExecutor creates DurableAgentThread instances."""
+        executor = ClientAgentExecutor(mock_client)
+
+        thread = executor.get_new_thread("test_agent")
+
+        assert isinstance(thread, DurableAgentThread)
+
+    def test_client_executor_forwards_kwargs_to_thread(self, mock_client: Mock) -> None:
+        """Verify ClientAgentExecutor forwards kwargs to DurableAgentThread creation."""
+        executor = ClientAgentExecutor(mock_client)
+
+        thread = executor.get_new_thread("test_agent", service_thread_id="client-123")
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.service_thread_id == "client-123"
+
+    def test_orchestration_executor_creates_durable_thread(
+        self, orchestration_executor: OrchestrationAgentExecutor
+    ) -> None:
+        """Verify OrchestrationAgentExecutor creates DurableAgentThread instances."""
+        thread = orchestration_executor.get_new_thread("test_agent")
+
+        assert isinstance(thread, DurableAgentThread)
+
+    def test_orchestration_executor_forwards_kwargs_to_thread(
+        self, orchestration_executor: OrchestrationAgentExecutor
+    ) -> None:
+        """Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentThread creation."""
+        thread = orchestration_executor.get_new_thread("test_agent", service_thread_id="orch-456")
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.service_thread_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 AgentRunResponse (synchronous)."""
+        result = client_executor.run_durable_agent("test_agent", sample_run_request)
+
+        # Verify it returns an AgentRunResponse (synchronous, not a coroutine)
+        assert isinstance(result, AgentRunResponse)
+        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 AgentRunResponse (should timeout after 2 attempts)
+        assert isinstance(result, AgentRunResponse)
+
+        # 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) -> 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)
+
+        # Measure time taken
+        start = time.time()
+        result = executor.run_durable_agent("test_agent", sample_run_request)
+        elapsed = time.time() - start
+
+        # Should take roughly 3 * 0.01 = 0.03 seconds (plus overhead)
+        # Be generous with timing to avoid flakiness
+        assert elapsed < 0.2  # Should be quick with 0.01 interval
+        assert isinstance(result, AgentRunResponse)
+
+
+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_thread_session_id(
+        self,
+        mock_orchestration_context: Mock,
+        orchestration_executor: OrchestrationAgentExecutor,
+        sample_run_request: RunRequest,
+    ) -> None:
+        """Verify executor uses thread's session ID when provided."""
+        # Create thread with specific session ID
+        session_id = AgentSessionId(name="test_agent", key="specific-key-123")
+        thread = DurableAgentThread.from_session_id(session_id)
+
+        result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, thread=thread)
+
+        # 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, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
+    ) -> None:
+        """Verify DurableAgentTask converts successful entity result to AgentRunResponse."""
+        mock_entity_task.is_failed = False
+        mock_entity_task.get_result = Mock(return_value=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, AgentRunResponse)
+        assert len(result.messages) == 1
+        assert result.messages[0].role == Role.ASSISTANT
+
+    def test_durable_agent_task_propagates_failure(self, mock_entity_task: Mock) -> None:
+        """Verify DurableAgentTask propagates task failures."""
+        mock_entity_task.is_failed = True
+        mock_entity_task.get_exception = Mock(return_value=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
+        exception = task.get_exception()
+        assert isinstance(exception, ValueError)
+        assert str(exception) == "Entity error"
+
+    def test_durable_agent_task_validates_response_format(self, mock_entity_task: Mock) -> None:
+        """Verify DurableAgentTask validates response format when provided."""
+        mock_entity_task.is_failed = False
+        mock_entity_task.get_result = Mock(
+            return_value={
+                "messages": [{"role": "assistant", "contents": [{"type": "text", "text": '{"answer": "42"}'}]}],
+                "created_at": "2025-12-30T10:00:00Z",
+            }
+        )
+
+        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, AgentRunResponse)
+
+    def test_durable_agent_task_ignores_duplicate_completion(
+        self, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
+    ) -> None:
+        """Verify DurableAgentTask ignores duplicate completion calls."""
+        mock_entity_task.is_failed = False
+        mock_entity_task.get_result = Mock(return_value=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
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_models.py

@@ -18,9 +18,10 @@ class TestRunRequest:
 
     def test_init_with_defaults(self) -> None:
         """Test RunRequest initialization with defaults."""
-        request = RunRequest(message="Hello")
+        request = RunRequest(message="Hello", correlation_id="corr-001")
 
         assert request.message == "Hello"
+        assert request.correlation_id == "corr-001"
         assert request.role == Role.USER
         assert request.response_format is None
         assert request.enable_tool_calls is True
@@ -30,30 +31,33 @@ class TestRunRequest:
         schema = ModuleStructuredResponse
         request = RunRequest(
             message="Hello",
+            correlation_id="corr-002",
             role=Role.SYSTEM,
             response_format=schema,
             enable_tool_calls=False,
         )
 
         assert request.message == "Hello"
+        assert request.correlation_id == "corr-002"
         assert request.role == Role.SYSTEM
         assert request.response_format is schema
         assert request.enable_tool_calls is False
 
     def test_init_coerces_string_role(self) -> None:
         """Ensure string role values are coerced into Role instances."""
-        request = RunRequest(message="Hello", role="system")  # type: ignore[arg-type]
+        request = RunRequest(message="Hello", correlation_id="corr-003", role="system")  # type: ignore[arg-type]
 
         assert request.role == Role.SYSTEM
 
     def test_to_dict_with_defaults(self) -> None:
         """Test to_dict with default values."""
-        request = RunRequest(message="Test message")
+        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["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
 
@@ -62,6 +66,7 @@ class TestRunRequest:
         schema = ModuleStructuredResponse
         request = RunRequest(
             message="Hello",
+            correlation_id="corr-005",
             role=Role.ASSISTANT,
             response_format=schema,
             enable_tool_calls=False,
@@ -69,6 +74,7 @@ class TestRunRequest:
         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__
@@ -78,16 +84,17 @@ class TestRunRequest:
 
     def test_from_dict_with_defaults(self) -> None:
         """Test from_dict with minimal data."""
-        data = {"message": "Hello"}
+        data = {"message": "Hello", "correlationId": "corr-006"}
         request = RunRequest.from_dict(data)
 
         assert request.message == "Hello"
+        assert request.correlation_id == "corr-006"
         assert request.role == Role.USER
         assert request.enable_tool_calls 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", "thread_id": "ignored"})
+        request = RunRequest.from_dict({"message": "Hello", "correlationId": "corr-007", "thread_id": "ignored"})
 
         assert request.message == "Hello"
 
@@ -95,6 +102,7 @@ class TestRunRequest:
         """Test from_dict with all fields."""
         data = {
             "message": "Test",
+            "correlationId": "corr-008",
             "role": "system",
             "response_format": {
                 "__response_schema_type__": "pydantic_model",
@@ -106,13 +114,14 @@ class TestRunRequest:
         request = RunRequest.from_dict(data)
 
         assert request.message == "Test"
+        assert request.correlation_id == "corr-008"
         assert request.role == Role.SYSTEM
         assert request.response_format is ModuleStructuredResponse
         assert request.enable_tool_calls is False
 
-    def test_from_dict_with_unknown_role_preserves_value(self) -> None:
+    def test_from_dict_unknown_role_preserves_value(self) -> None:
         """Test from_dict keeps custom roles intact."""
-        data = {"message": "Test", "role": "reviewer"}
+        data = {"message": "Test", "correlationId": "corr-009", "role": "reviewer"}
         request = RunRequest.from_dict(data)
 
         assert request.role.value == "reviewer"
@@ -120,15 +129,22 @@ class TestRunRequest:
 
     def test_from_dict_empty_message(self) -> None:
         """Test from_dict with empty message."""
-        request = RunRequest.from_dict({})
+        request = RunRequest.from_dict({"correlationId": "corr-010"})
 
         assert request.message == ""
+        assert request.correlation_id == "corr-010"
         assert request.role == 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=Role.SYSTEM,
             response_format=ModuleStructuredResponse,
             enable_tool_calls=False,
@@ -138,6 +154,7 @@ class TestRunRequest:
         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
@@ -146,6 +163,7 @@ class TestRunRequest:
         """Ensure Pydantic response formats serialize and deserialize properly."""
         original = RunRequest(
             message="Structured",
+            correlation_id="corr-012",
             response_format=ModuleStructuredResponse,
         )
 
@@ -186,7 +204,7 @@ class TestRunRequest:
         original = RunRequest(
             message="Test message",
             role=Role.SYSTEM,
-            correlation_id="corr-123",
+            correlation_id="corr-124",
         )
 
         data = original.to_dict()
@@ -200,6 +218,7 @@ class TestRunRequest:
         """Test RunRequest initialization with orchestration_id."""
         request = RunRequest(
             message="Test message",
+            correlation_id="corr-125",
             orchestration_id="orch-123",
         )
 
@@ -210,6 +229,7 @@ class TestRunRequest:
         """Test to_dict includes orchestrationId."""
         request = RunRequest(
             message="Test",
+            correlation_id="corr-126",
             orchestration_id="orch-456",
         )
         data = request.to_dict()
@@ -221,6 +241,7 @@ class TestRunRequest:
         """Test to_dict excludes orchestrationId when not set."""
         request = RunRequest(
             message="Test",
+            correlation_id="corr-127",
         )
         data = request.to_dict()
 
@@ -230,6 +251,7 @@ class TestRunRequest:
         """Test from_dict with orchestrationId."""
         data = {
             "message": "Test",
+            "correlationId": "corr-128",
             "orchestrationId": "orch-789",
         }
         request = RunRequest.from_dict(data)
@@ -242,7 +264,7 @@ class TestRunRequest:
         original = RunRequest(
             message="Test message",
             role=Role.SYSTEM,
-            correlation_id="corr-123",
+            correlation_id="corr-129",
             orchestration_id="orch-123",
         )
 

python/packages/durabletask/tests/test_orchestration_context.py

@@ -0,0 +1,98 @@
+# 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 AgentProtocol
+
+from agent_framework_durabletask import DurableAgentThread
+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, AgentProtocol)
+
+    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_threads(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
+        """Verify agent from context can create DurableAgentThread instances."""
+        agent = agent_context.get_agent("assistant")
+
+        thread = agent.get_new_thread()
+
+        assert isinstance(thread, DurableAgentThread)
+
+    def test_orchestration_agent_thread_with_parameters(
+        self, agent_context: DurableAIAgentOrchestrationContext
+    ) -> None:
+        """Verify agent can create threads with custom parameters."""
+        agent = agent_context.get_agent("assistant")
+
+        thread = agent.get_new_thread(service_thread_id="orch-session-456")
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.service_thread_id == "orch-session-456"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

python/packages/durabletask/tests/test_shim.py

@@ -0,0 +1,206 @@
+# 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
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import AgentProtocol, ChatMessage
+from pydantic import BaseModel
+
+from agent_framework_durabletask import DurableAgentThread
+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_thread = Mock(return_value=DurableAgentThread())
+
+    # Mock get_run_request to create actual RunRequest objects
+    def create_run_request(
+        message: str, response_format: type[BaseModel] | None = None, enable_tool_calls: bool = True
+    ) -> RunRequest:
+        import uuid
+
+        return RunRequest(
+            message=message,
+            correlation_id=str(uuid.uuid4()),
+            response_format=response_format,
+            enable_tool_calls=enable_tool_calls,
+        )
+
+    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 ChatMessage objects."""
+        chat_msg = ChatMessage(role="user", text="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 ChatMessage objects."""
+        messages = [
+            ChatMessage(role="user", text="Message 1"),
+            ChatMessage(role="assistant", text="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_thread_parameter(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify run forwards thread parameter to executor."""
+        thread = DurableAgentThread(service_thread_id="test-thread")
+        test_agent.run("message", thread=thread)
+
+        mock_executor.run_durable_agent.assert_called_once()
+        _, kwargs = mock_executor.run_durable_agent.call_args
+        assert kwargs["thread"] == thread
+
+    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", 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 TestDurableAIAgentProtocolCompliance:
+    """Test that DurableAIAgent implements AgentProtocol correctly."""
+
+    def test_agent_implements_protocol(self, test_agent: DurableAIAgent[Any]) -> None:
+        """Verify DurableAIAgent implements AgentProtocol."""
+        assert isinstance(test_agent, AgentProtocol)
+
+    def test_agent_has_required_properties(self, test_agent: DurableAIAgent[Any]) -> None:
+        """Verify DurableAIAgent has all required AgentProtocol 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 TestDurableAIAgentThreadManagement:
+    """Test thread creation and management."""
+
+    def test_get_new_thread_delegates_to_executor(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify get_new_thread delegates to executor."""
+        mock_thread = DurableAgentThread()
+        mock_executor.get_new_thread.return_value = mock_thread
+
+        thread = test_agent.get_new_thread()
+
+        mock_executor.get_new_thread.assert_called_once_with("test_agent")
+        assert thread == mock_thread
+
+    def test_get_new_thread_forwards_kwargs(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
+        """Verify get_new_thread forwards kwargs to executor."""
+        mock_thread = DurableAgentThread(service_thread_id="thread-123")
+        mock_executor.get_new_thread.return_value = mock_thread
+
+        test_agent.get_new_thread(service_thread_id="thread-123")
+
+        mock_executor.get_new_thread.assert_called_once()
+        _, kwargs = mock_executor.get_new_thread.call_args
+        assert kwargs["service_thread_id"] == "thread-123"
+
+
+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_worker.py

@@ -0,0 +1,168 @@
+# 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
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])