diff --git a/src/bonsai/bonsai/bim/module/document/operator.py b/src/bonsai/bonsai/bim/module/document/operator.py
index 2ed1a0a8d5..da75e16e63 100644
--- a/src/bonsai/bonsai/bim/module/document/operator.py
+++ b/src/bonsai/bonsai/bim/module/document/operator.py
@@ -132,13 +132,14 @@ class RemoveDocument(bpy.types.Operator, tool.Ifc.Operator):
 class AssignDocument(bpy.types.Operator, tool.Ifc.Operator):
     bl_idname = "bim.assign_document"
     bl_label = "Assign Document"
+    bl_description = "Assign active document to the selected objects."
     bl_options = {"REGISTER", "UNDO"}
     obj: bpy.props.StringProperty()
     document: bpy.props.IntProperty()
 
     def _execute(self, context):
         document = tool.Ifc.get().by_id(self.document)
-        objs = [bpy.data.objects.get(self.obj)] if self.obj else tool.Blender.get_selected_objects()
+        objs = [bpy.data.objects[self.obj]] if self.obj else tool.Blender.get_selected_objects()
         for obj in objs:
             element = tool.Ifc.get_entity(obj)
             if element:
diff --git a/src/bonsai/bonsai/bim/module/document/prop.py b/src/bonsai/bonsai/bim/module/document/prop.py
index ef4fb29ab9..f4e4b16686 100644
--- a/src/bonsai/bonsai/bim/module/document/prop.py
+++ b/src/bonsai/bonsai/bim/module/document/prop.py
@@ -30,7 +30,7 @@ from bpy.props import (
     FloatVectorProperty,
     CollectionProperty,
 )
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 
 def update_document_name(self: "Document", context: bpy.types.Context) -> None:
@@ -79,3 +79,7 @@ class BIMDocumentProperties(PropertyGroup):
         breadcrumbs: bpy.types.bpy_prop_collection_idprop[StrProperty]
         active_document_index: int
         is_editing: bool
+
+    @property
+    def active_document(self) -> Union[Document, None]:
+        return tool.Blender.get_active_uilist_element(self.documents, self.active_document_index)
diff --git a/src/bonsai/bonsai/bim/module/document/ui.py b/src/bonsai/bonsai/bim/module/document/ui.py
index fefb87fe6c..5897edbad8 100644
--- a/src/bonsai/bonsai/bim/module/document/ui.py
+++ b/src/bonsai/bonsai/bim/module/document/ui.py
@@ -61,14 +61,17 @@ class BIM_PT_documents(Panel):
         if self.props.breadcrumbs:
             row.operator("bim.add_document_reference", text="", icon="FILE_HIDDEN")
 
+        active_document = self.props.active_document
+
         if self.props.active_document_id:
             row.operator("bim.edit_document", text="", icon="CHECKMARK")
             row.operator("bim.disable_editing_document", text="", icon="CANCEL")
-        elif self.props.documents and self.props.active_document_index < len(self.props.documents):
-            ifc_definition_id = self.props.documents[self.props.active_document_index].ifc_definition_id
+        elif active_document:
+            ifc_definition_id = active_document.ifc_definition_id
             row.operator("bim.select_document_objects", text="", icon="RESTRICT_SELECT_OFF").document = (
                 ifc_definition_id
             )
+            row.operator("bim.assign_document", text="", icon="BRUSH_DATA").document = ifc_definition_id
             row.operator("bim.enable_editing_document", text="", icon="GREASEPENCIL").document = ifc_definition_id
             row.operator("bim.remove_document", text="", icon="X").document = ifc_definition_id