Python: Add orchestration ID to durable agent entity state and code refactor (#2484)

* Initial plan

* Add orchestration ID to durable agent entity state for Python

Co-authored-by: larohra <41490930+larohra@users.noreply.github.com>

* Fix type safety checks

* Fix tests

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: larohra <41490930+larohra@users.noreply.github.com>
Co-authored-by: Laveesh Rohra <larohra@microsoft.com>

python/packages/azurefunctions/tests/test_entities.py

@@ -79,7 +79,7 @@ class TestAgentEntityInit:
         assert entity.agent == mock_agent
         assert len(entity.state.data.conversation_history) == 0
         assert entity.state.data.extension_data is None
-        assert entity.state.schema_version == "1.0.0"
+        assert entity.state.schema_version == DurableAgentState.SCHEMA_VERSION
 
     def test_init_stores_agent_reference(self) -> None:
         """Test that the agent reference is stored correctly."""
@@ -124,8 +124,7 @@ class TestAgentEntityRunAgent:
         # Verify agent.run was called
         mock_agent.run.assert_called_once()
         _, kwargs = mock_agent.run.call_args
-        sent_messages = kwargs.get("messages")
-        assert isinstance(sent_messages, list)
+        sent_messages: list[Any] = kwargs.get("messages")
         assert len(sent_messages) == 1
         sent_message = sent_messages[0]
         assert isinstance(sent_message, ChatMessage)
@@ -910,5 +909,98 @@ class TestRunRequestSupport:
         assert text_found, f"Response text not found in message: {message}"
 
 
+class TestDurableAgentStateRequestOrchestrationId:
+    """Test suite for DurableAgentStateRequest orchestration_id field."""
+
+    def test_request_with_orchestration_id(self) -> None:
+        """Test creating a request with an orchestration_id."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+            orchestration_id="orch-456",
+        )
+
+        assert request.orchestration_id == "orch-456"
+
+    def test_request_to_dict_includes_orchestration_id(self) -> None:
+        """Test that to_dict includes orchestrationId when set."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+            orchestration_id="orch-789",
+        )
+
+        data = request.to_dict()
+
+        assert "orchestrationId" in data
+        assert data["orchestrationId"] == "orch-789"
+
+    def test_request_to_dict_excludes_orchestration_id_when_none(self) -> None:
+        """Test that to_dict excludes orchestrationId when not set."""
+        request = DurableAgentStateRequest(
+            correlation_id="corr-123",
+            created_at=datetime.now(),
+            messages=[
+                DurableAgentStateMessage(
+                    role="user",
+                    contents=[DurableAgentStateTextContent(text="test")],
+                )
+            ],
+        )
+
+        data = request.to_dict()
+
+        assert "orchestrationId" not in data
+
+    def test_request_from_dict_with_orchestration_id(self) -> None:
+        """Test from_dict correctly parses orchestrationId."""
+        data = {
+            "$type": "request",
+            "correlationId": "corr-123",
+            "createdAt": "2024-01-01T00:00:00Z",
+            "messages": [{"role": "user", "contents": [{"$type": "text", "text": "test"}]}],
+            "orchestrationId": "orch-from-dict",
+        }
+
+        request = DurableAgentStateRequest.from_dict(data)
+
+        assert request.orchestration_id == "orch-from-dict"
+
+    def test_request_from_run_request_with_orchestration_id(self) -> None:
+        """Test from_run_request correctly transfers orchestration_id."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+            orchestration_id="orch-from-run-request",
+        )
+
+        durable_request = DurableAgentStateRequest.from_run_request(run_request)
+
+        assert durable_request.orchestration_id == "orch-from-run-request"
+
+    def test_request_from_run_request_without_orchestration_id(self) -> None:
+        """Test from_run_request correctly handles missing orchestration_id."""
+        run_request = RunRequest(
+            message="test message",
+            correlation_id="corr-run",
+        )
+
+        durable_request = DurableAgentStateRequest.from_run_request(run_request)
+
+        assert durable_request.orchestration_id is None
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "--tb=short"])

python/packages/azurefunctions/tests/test_models.py

@@ -336,6 +336,71 @@ class TestRunRequest:
         assert restored.correlation_id == original.correlation_id
         assert restored.thread_id == original.thread_id
 
+    def test_init_with_orchestration_id(self) -> None:
+        """Test RunRequest initialization with orchestration_id."""
+        request = RunRequest(
+            message="Test message",
+            thread_id="thread-orch-init",
+            orchestration_id="orch-123",
+        )
+
+        assert request.message == "Test message"
+        assert request.orchestration_id == "orch-123"
+
+    def test_to_dict_with_orchestration_id(self) -> None:
+        """Test to_dict includes orchestrationId."""
+        request = RunRequest(
+            message="Test",
+            thread_id="thread-orch-to-dict",
+            orchestration_id="orch-456",
+        )
+        data = request.to_dict()
+
+        assert data["message"] == "Test"
+        assert data["orchestrationId"] == "orch-456"
+
+    def test_to_dict_excludes_orchestration_id_when_none(self) -> None:
+        """Test to_dict excludes orchestrationId when not set."""
+        request = RunRequest(
+            message="Test",
+            thread_id="thread-orch-none",
+        )
+        data = request.to_dict()
+
+        assert "orchestrationId" not in data
+
+    def test_from_dict_with_orchestration_id(self) -> None:
+        """Test from_dict with orchestrationId."""
+        data = {
+            "message": "Test",
+            "orchestrationId": "orch-789",
+            "thread_id": "thread-orch-from-dict",
+        }
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Test"
+        assert request.orchestration_id == "orch-789"
+        assert request.thread_id == "thread-orch-from-dict"
+
+    def test_round_trip_with_orchestration_id(self) -> None:
+        """Test round-trip to_dict and from_dict with orchestration_id."""
+        original = RunRequest(
+            message="Test message",
+            thread_id="thread-123",
+            role=Role.SYSTEM,
+            correlation_id="corr-123",
+            orchestration_id="orch-123",
+        )
+
+        data = original.to_dict()
+        restored = RunRequest.from_dict(data)
+
+        assert restored.message == original.message
+        assert restored.role == original.role
+        assert restored.correlation_id == original.correlation_id
+        assert restored.orchestration_id == original.orchestration_id
+        assert restored.thread_id == original.thread_id
+
 
 class TestModelIntegration:
     """Test suite for integration between models."""

python/packages/azurefunctions/tests/test_orchestration.py

@@ -302,6 +302,28 @@ class TestDurableAIAgent:
         assert request["correlationId"] == "correlation-guid"
         assert "thread_id" in request
         assert request["thread_id"] == "thread-guid"
+        # Verify orchestration ID is set from context.instance_id
+        assert "orchestrationId" in request
+        assert request["orchestrationId"] == "test-instance-001"
+
+    def test_run_sets_orchestration_id(self) -> None:
+        """Test that run() sets the orchestration_id from context.instance_id."""
+        mock_context = Mock()
+        mock_context.instance_id = "my-orchestration-123"
+        mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
+
+        entity_task = _create_entity_task()
+        mock_context.call_entity = Mock(return_value=entity_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+        thread = agent.get_new_thread()
+
+        agent.run(messages="Test", thread=thread)
+
+        call_args = mock_context.call_entity.call_args
+        request = call_args[0][2]
+
+        assert request["orchestrationId"] == "my-orchestration-123"
 
     def test_run_without_thread(self) -> None:
         """Test that run() works without explicit thread (creates unique session key)."""