Add support for enabled features in config

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -69,7 +69,9 @@ def start(self, event_scheduler):
         send_heartbeats_every_x_secs(self, self.heartbeat_secs, event_scheduler)
         start_polling_for_changes(self, event_scheduler)
 
-        if is_feature_enabled("sse"):
+        if is_feature_enabled("sse") or self.conf.is_feature_enabled(
+            "realtime_updates"
+        ):
             listen_for_config_updates(self, event_scheduler)
 
     def report_initial_stats(self):

aikido_zen/background_process/cloud_connection_manager/init_test.py

@@ -1,4 +1,5 @@
 import pytest
+from unittest.mock import patch, MagicMock
 from aikido_zen.helpers.token import Token
 from aikido_zen.background_process.api.http_api import ReportingApiHTTP
 from aikido_zen.background_process.service_config import ServiceConfig
@@ -45,4 +46,112 @@ def test_cloud_connection_manager_empty_serverless():
         CloudConnectionManager(block, api, token, serverless)
 
 
+@patch("aikido_zen.background_process.cloud_connection_manager.on_start")
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates"
+)
+def test_start_enables_sse_when_server_flag_set(
+    mock_listen_for_config_updates,
+    mock_send_heartbeats,
+    mock_start_polling,
+    mock_on_start,
+    setup_cloud_connection_manager,
+):
+    """SSE listening should start when the server enables the realtime_updates feature flag"""
+    manager = setup_cloud_connection_manager
+    mock_on_start.return_value = {"success": True}
+    manager.conf.update_enabled_features(["realtime_updates"])
+
+    manager.start(event_scheduler=MagicMock())
+
+    mock_listen_for_config_updates.assert_called_once()
+
+
+@patch("aikido_zen.background_process.cloud_connection_manager.on_start")
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates"
+)
+def test_start_does_not_enable_sse_without_flag(
+    mock_listen_for_config_updates,
+    mock_send_heartbeats,
+    mock_start_polling,
+    mock_on_start,
+    setup_cloud_connection_manager,
+):
+    """SSE listening should not start when neither the env var nor the server flag is set"""
+    manager = setup_cloud_connection_manager
+    mock_on_start.return_value = {"success": True}
+
+    manager.start(event_scheduler=MagicMock())
+
+    mock_listen_for_config_updates.assert_not_called()
+
+
+@patch("aikido_zen.background_process.cloud_connection_manager.on_start")
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates"
+)
+def test_start_enables_sse_via_env_var(
+    mock_listen_for_config_updates,
+    mock_send_heartbeats,
+    mock_start_polling,
+    mock_on_start,
+    setup_cloud_connection_manager,
+    monkeypatch,
+):
+    """SSE listening should still start via the AIKIDO_FEATURE_SSE env var, without the server flag"""
+    monkeypatch.setenv("AIKIDO_FEATURE_SSE", "true")
+    manager = setup_cloud_connection_manager
+    mock_on_start.return_value = {"success": True}
+
+    manager.start(event_scheduler=MagicMock())
+
+    mock_listen_for_config_updates.assert_called_once()
+
+
+@patch("aikido_zen.background_process.cloud_connection_manager.on_start")
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.start_polling_for_changes"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.send_heartbeats_every_x_secs"
+)
+@patch(
+    "aikido_zen.background_process.cloud_connection_manager.listen_for_config_updates"
+)
+def test_start_does_not_enable_sse_for_unrelated_flag(
+    mock_listen_for_config_updates,
+    mock_send_heartbeats,
+    mock_start_polling,
+    mock_on_start,
+    setup_cloud_connection_manager,
+):
+    """A server feature flag other than realtime_updates should not enable SSE"""
+    manager = setup_cloud_connection_manager
+    mock_on_start.return_value = {"success": True}
+    manager.conf.update_enabled_features(["some_other_feature"])
+
+    manager.start(event_scheduler=MagicMock())
+
+    mock_listen_for_config_updates.assert_not_called()
+
+
 # Additional tests can be added here for other edge cases or scenarios

aikido_zen/background_process/cloud_connection_manager/update_service_config.py

@@ -38,3 +38,6 @@ def update_service_config(connection_manager, res):
 
     if "domains" in res:
         connection_manager.conf.update_outbound_domains(res["domains"])
+
+    if isinstance(res.get("enabledFeatures"), list):
+        connection_manager.conf.update_enabled_features(res["enabledFeatures"])

aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py

@@ -260,6 +260,97 @@ def test_update_service_config_excluded_user_ids_from_rate_limiting():
     assert connection_manager.conf.is_user_excluded_from_rate_limiting("user3") is False
 
 
+def test_update_service_config_enabled_features():
+    """Test that update_service_config handles enabledFeatures"""
+    connection_manager = MagicMock()
+    connection_manager.conf = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+    connection_manager.block = False
+
+    res = {
+        "success": True,
+        "enabledFeatures": ["realtime_updates"],
+    }
+
+    update_service_config(connection_manager, res)
+
+    assert connection_manager.conf.is_feature_enabled("realtime_updates") is True
+    assert connection_manager.conf.is_feature_enabled("some_other_feature") is False
+
+
+def test_update_service_config_enabled_features_not_array():
+    """Test that update_service_config ignores non-array enabledFeatures"""
+    connection_manager = MagicMock()
+    connection_manager.conf = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+    connection_manager.block = False
+
+    res = {
+        "success": True,
+        "enabledFeatures": "not-an-array",
+    }
+
+    update_service_config(connection_manager, res)
+
+    assert connection_manager.conf.is_feature_enabled("not-an-array") is False
+
+
+def test_update_service_config_enabled_features_empty_list():
+    """Test that update_service_config clears enabled features with an empty list"""
+    connection_manager = MagicMock()
+    connection_manager.conf = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+    connection_manager.block = False
+    connection_manager.conf.update_enabled_features(["realtime_updates"])
+
+    res = {
+        "success": True,
+        "enabledFeatures": [],
+    }
+
+    update_service_config(connection_manager, res)
+
+    assert connection_manager.conf.is_feature_enabled("realtime_updates") is False
+
+
+def test_update_service_config_enabled_features_not_updated_on_failure():
+    """Test that enabledFeatures is left untouched when the response indicates failure"""
+    connection_manager = MagicMock()
+    connection_manager.conf = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+    connection_manager.block = False
+    connection_manager.conf.update_enabled_features(["realtime_updates"])
+
+    res = {
+        "success": False,
+        "enabledFeatures": [],
+    }
+
+    update_service_config(connection_manager, res)
+
+    assert connection_manager.conf.is_feature_enabled("realtime_updates") is True
+
+
 def test_update_service_config_excluded_user_ids_not_array():
     """Test that update_service_config ignores non-array excludedUserIdsFromRateLimiting"""
     connection_manager = MagicMock()

aikido_zen/background_process/service_config.py

@@ -25,6 +25,7 @@ def __init__(
         self.block_new_outgoing_requests = False
         self.outbound_domains = {}
         self.excluded_user_ids_from_rate_limiting = set()
+        self.enabled_features = set()
 
     def update(
         self,
@@ -84,6 +85,14 @@ def is_user_excluded_from_rate_limiting(self, user_id):
         """Checks if the user ID is excluded from rate limiting"""
         return str(user_id) in self.excluded_user_ids_from_rate_limiting
 
+    def update_enabled_features(self, features):
+        """Sets the server-provided feature flags"""
+        self.enabled_features = set(features)
+
+    def is_feature_enabled(self, feature) -> bool:
+        """Checks if `feature` is enabled via server-provided config"""
+        return feature in self.enabled_features
+
     def update_outbound_domains(self, domains):
         self.outbound_domains = {
             domain["hostname"]: domain["mode"] for domain in domains

aikido_zen/background_process/service_config_test.py

@@ -321,6 +321,28 @@ def test_service_config_with_empty_allowlist():
     assert len(admin_endpoint["allowedIPAddresses"]) == 0
 
 
+def test_service_config_enabled_features():
+    """Test the update_enabled_features and is_feature_enabled methods"""
+    config = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+
+    # Initially empty
+    assert config.is_feature_enabled("realtime_updates") is False
+
+    config.update_enabled_features(["realtime_updates"])
+    assert config.is_feature_enabled("realtime_updates") is True
+    assert config.is_feature_enabled("some_other_feature") is False
+
+    # Update replaces the set
+    config.update_enabled_features([])
+    assert config.is_feature_enabled("realtime_updates") is False
+
+
 def test_excluded_user_ids_from_rate_limiting():
     config = ServiceConfig(
         endpoints=[],