diff --git a/capellambse_context_diagrams/context.py b/capellambse_context_diagrams/context.py
index 828c3668..9635d0c5 100644
--- a/capellambse_context_diagrams/context.py
+++ b/capellambse_context_diagrams/context.py
@@ -364,6 +364,7 @@ def _create_diagram(self, params: dict[str, t.Any]) -> cdiagram.Diagram:
         )
         data, legend = tree_view.collector(self, params)
         params["elkdata"] = data
+        add_context(data)
         class_diagram = super()._create_diagram(params)
         width, height = class_diagram.viewport.size
         axis: t.Literal["x", "y"]
@@ -379,6 +380,21 @@ def _create_diagram(self, params: dict[str, t.Any]) -> cdiagram.Diagram:
         return class_diagram
 
 
+def add_context(data: _elkjs.ELKInputData) -> None:
+    ids: set[str] = set()
+
+    def get_ids(obj: _elkjs.ELKInputChild):
+        ids.add(obj["id"])
+        for cobj in obj.get("children", []):
+            get_ids(cobj)
+
+    for child in data["children"]:
+        get_ids(child)
+
+    for child in data["children"] + data["edges"]:
+        child["context"] = list(ids)
+
+
 def stack_diagrams(
     first: cdiagram.Diagram,
     second: cdiagram.Diagram,