diff --git a/capellambse_context_diagrams/context.py b/capellambse_context_diagrams/context.py
index b956411b..b1663bda 100644
--- a/capellambse_context_diagrams/context.py
+++ b/capellambse_context_diagrams/context.py
@@ -448,15 +448,15 @@ def add_context(data: _elkjs.ELKOutputData, is_legend: bool = False) -> None:
 
     ids: set[str] = set()
 
-    def get_ids(obj: _elkjs.ELKOutputNode | _elkjs.ELKOutputEdge) -> None:
+    def get_ids(obj: _elkjs.ELKOutputElement) -> None:
         if obj["id"] and not obj["id"].startswith("g_"):
             ids.add(obj["id"])
-        for cobj in obj.get("children", []):
-            if cobj["type"] == "node":
+        for cobj in obj.get("children", []):  # type: ignore[attr-defined]
+            if cobj["type"] in {"node", "port", "junction", "edge"}:
                 get_ids(cobj)
 
     for child in data["children"]:
-        if child["type"] == "node" or child["type"] == "edge":
+        if child["type"] in {"node", "port", "junction", "edge"}:
             get_ids(child)
 
     for child in data["children"]: