diff --git a/ci/optimized_targets.py b/ci/optimized_targets.py
index 8a529c7..116d6f8 100644
--- a/ci/optimized_targets.py
+++ b/ci/optimized_targets.py
@@ -14,9 +14,10 @@
 # limitations under the License.
 
 from abc import ABC
-from typing import Self
 import argparse
 import functools
+from typing import Self
+from build_context import BuildContext
 
 
 class OptimizedBuildTarget(ABC):
@@ -30,7 +31,7 @@
   def __init__(
       self,
       target: str,
-      build_context: dict[str, any],
+      build_context: BuildContext,
       args: argparse.Namespace,
   ):
     self.target = target
@@ -38,13 +39,13 @@
     self.args = args
 
   def get_build_targets(self) -> set[str]:
-    features = self.build_context.get('enabledBuildFeatures', [])
+    features = self.build_context.enabled_build_features
     if self.get_enabled_flag() in features:
       return self.get_build_targets_impl()
     return {self.target}
 
   def package_outputs(self):
-    features = self.build_context.get('enabledBuildFeatures', [])
+    features = self.build_context.enabled_build_features
     if self.get_enabled_flag() in features:
       return self.package_outputs_impl()