diff --git a/src/alpine_bits_python/alpinebits_server.py b/src/alpine_bits_python/alpinebits_server.py
index a9cfa17..55f43dc 100644
--- a/src/alpine_bits_python/alpinebits_server.py
+++ b/src/alpine_bits_python/alpinebits_server.py
@@ -196,7 +196,7 @@ class ServerCapabilities:
     """
 
     def __init__(self):
-        self.action_registry: Dict[str, Type[AlpineBitsAction]] = {}
+        self.action_registry: Dict[AlpineBitsActionName, Type[AlpineBitsAction]] = {}
         self._discover_actions()
         self.capability_dict = None
 
@@ -214,8 +214,8 @@ class ServerCapabilities:
                 if self._is_action_implemented(obj):
                     action_instance = obj()
                     if hasattr(action_instance, "name"):
-                        # Use capability name for the registry key
-                        self.action_registry[action_instance.name.capability_name] = obj
+                        # Use capability attribute as registry key
+                        self.action_registry[action_instance.name] = obj
 
     def _is_action_implemented(self, action_class: Type[AlpineBitsAction]) -> bool:
         """
@@ -234,7 +234,7 @@ class ServerCapabilities:
         """
         versions_dict = {}
 
-        for action_name, action_class in self.action_registry.items():
+        for action_enum, action_class in self.action_registry.items():
             action_instance = action_class()
 
             # Get supported versions for this action
@@ -250,7 +250,7 @@ class ServerCapabilities:
                 if version_str not in versions_dict:
                     versions_dict[version_str] = {"version": version_str, "actions": []}
 
-                action_dict = {"action": action_name}
+                action_dict = {"action": action_enum.capability_name}
 
                 # Add supports field if the action has custom supports
                 if hasattr(action_instance, "supports") and action_instance.supports:
@@ -714,14 +714,14 @@ class AlpineBitsServer:
             )
 
         # Check if we have an implementation for this action
-        capability_name = action_enum.capability_name
-        if capability_name not in self._action_instances:
+        
+        if action_enum not in self._action_instances:
             return AlpineBitsResponse(
                 f"Error: Action {request_action_name} is not implemented",
                 HttpStatusCode.BAD_REQUEST,
             )
 
-        action_instance: AlpineBitsAction = self._action_instances[capability_name]
+        action_instance: AlpineBitsAction = self._action_instances[action_enum]
 
         # Check if the action supports the requested version
         if not await action_instance.check_version_supported(version_enum):
@@ -733,7 +733,7 @@ class AlpineBitsServer:
         # Handle the request
         try:
             # Special case for ping action - pass server capabilities
-            if capability_name == "action_OTA_Ping":
+            if action_enum == AlpineBitsActionName.OTA_PING:
                 return await action_instance.handle(
                     action=request_action_name, request_xml=request_xml, version=version_enum, server_capabilities=self.capabilities, client_info=client_info
                 )