Support gradual rollout flag for realtime updates

agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java

@@ -59,8 +59,10 @@ public void run() {
         // one time check to report initial stats
         scheduler.schedule(new HeartbeatTask(api, true), 60, TimeUnit.SECONDS);
 
-        if (token != null && FeatureFlags.AIKIDO_FEATURE_SSE.isEnabled()) {
-            new RealtimeSSETask(new RealtimeSSEAPI(token), api).start();
+        if (token != null) {
+            if (FeatureFlags.AIKIDO_FEATURE_SSE.isEnabled() || ServiceConfigStore.isRealtimeUpdatesEnabled()) {
+                new RealtimeSSETask(new RealtimeSSEAPI(token), api).start();
+            }
         }
     }
 }

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java

@@ -16,6 +16,23 @@ public record APIResponse(
         List<Domain> domains,
         boolean receivedAnyStats,
         boolean block,
-        List<String> excludedUserIdsFromRateLimiting
+        List<String> excludedUserIdsFromRateLimiting,
+        List<String> enabledFeatures
 ) {
+    public APIResponse(
+            boolean success,
+            String error,
+            long configUpdatedAt,
+            List<Endpoint> endpoints,
+            List<String> blockedUserIds,
+            List<String> allowedIPAddresses,
+            boolean blockNewOutgoingRequests,
+            List<Domain> domains,
+            boolean receivedAnyStats,
+            boolean block,
+            List<String> excludedUserIdsFromRateLimiting
+    ) {
+        this(success, error, configUpdatedAt, endpoints, blockedUserIds, allowedIPAddresses, blockNewOutgoingRequests,
+                domains, receivedAnyStats, block, excludedUserIdsFromRateLimiting, List.of());
+    }
 }

agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfigStore.java

@@ -41,6 +41,15 @@ public static boolean isBlockedUserAgent(String userAgent) {
         }
     }
 
+    public static boolean isRealtimeUpdatesEnabled() {
+        mutex.readLock().lock();
+        try {
+            return config.isRealtimeUpdatesEnabled();
+        } finally {
+            mutex.readLock().unlock();
+        }
+    }
+
     public static void updateFromAPIResponse(APIResponse apiResponse) {
         mutex.writeLock().lock();
         try {

agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfiguration.java

@@ -25,6 +25,7 @@ public class ServiceConfiguration {
     private IPList bypassedIPs = new IPList();
     private HashSet<String> blockedUserIDs = new HashSet<>();
     private HashSet<String> excludedUserIdsFromRateLimiting = new HashSet<>();
+    private HashSet<String> enabledFeatures = new HashSet<>();
     private List<Endpoint> endpoints = new ArrayList<>();
     private OutboundDomains outboundDomains = new OutboundDomains();
 
@@ -47,6 +48,9 @@ public void updateConfig(APIResponse apiResponse) {
         if (apiResponse.excludedUserIdsFromRateLimiting() != null) {
             this.excludedUserIdsFromRateLimiting = new HashSet<>(apiResponse.excludedUserIdsFromRateLimiting());
         }
+        if (apiResponse.enabledFeatures() != null) {
+            this.enabledFeatures = new HashSet<>(apiResponse.enabledFeatures());
+        }
         if (apiResponse.endpoints() != null) {
             this.endpoints = apiResponse.endpoints();
         }
@@ -58,6 +62,10 @@ public boolean isBlockingEnabled() {
         return blockingEnabled;
     }
 
+    public boolean isRealtimeUpdatesEnabled() {
+        return enabledFeatures.contains("realtime_updates");
+    }
+
     public void setBlocking(boolean block) {
         this.blockingEnabled = block;
     }

agent_api/src/test/java/storage/ServiceConfigStoreTest.java

@@ -58,4 +58,18 @@ public void testShouldBlockOutgoingRequestUnknownWhenBlockNewEnabled() {
         ));
         assertTrue(ServiceConfigStore.shouldBlockOutgoingRequest("unknown.com"));
     }
+
+    @Test
+    public void testIsRealtimeUpdatesEnabledFalseByDefault() {
+        assertFalse(ServiceConfigStore.isRealtimeUpdatesEnabled());
+    }
+
+    @Test
+    public void testIsRealtimeUpdatesEnabledTrueWhenEnabled() {
+        ServiceConfigStore.updateFromAPIResponse(new APIResponse(
+                true, null, 0L, null, null, null,
+                false, null, true, false, null, List.of("realtime_updates")
+        ));
+        assertTrue(ServiceConfigStore.isRealtimeUpdatesEnabled());
+    }
 }

agent_api/src/test/java/storage/ServiceConfigurationTest.java

@@ -665,4 +665,58 @@ public void testIsIpBlockedWithAllowedIPsAndBlockedIPs() {
         assertTrue(resultBlocked.blocked());
         assertFalse(resultNotAllowedLocal.blocked());
     }
+
+    @Test
+    public void testIsRealtimeUpdatesEnabledDefaultsFalse() {
+        assertFalse(serviceConfiguration.isRealtimeUpdatesEnabled());
+    }
+
+    @Test
+    public void testUpdateConfigWithNullEnabledFeaturesDoesNotThrow() {
+        APIResponse apiResponse = new APIResponse(
+                true, null, 0L, null, null, null,
+                false, null, true, false, null, null
+        );
+        assertDoesNotThrow(() -> serviceConfiguration.updateConfig(apiResponse));
+        assertFalse(serviceConfiguration.isRealtimeUpdatesEnabled());
+    }
+
+    @Test
+    public void testIsRealtimeUpdatesEnabledTrueWhenPresent() {
+        APIResponse apiResponse = new APIResponse(
+                true, null, 0L, null, null, null,
+                false, null, true, false, null, List.of("realtime_updates")
+        );
+        serviceConfiguration.updateConfig(apiResponse);
+
+        assertTrue(serviceConfiguration.isRealtimeUpdatesEnabled());
+    }
+
+    @Test
+    public void testIsRealtimeUpdatesEnabledFalseWhenRemoved() {
+        APIResponse enabled = new APIResponse(
+                true, null, 0L, null, null, null,
+                false, null, true, false, null, List.of("realtime_updates")
+        );
+        serviceConfiguration.updateConfig(enabled);
+        assertTrue(serviceConfiguration.isRealtimeUpdatesEnabled());
+
+        APIResponse disabled = new APIResponse(
+                true, null, 0L, null, null, null,
+                false, null, true, false, null, List.of()
+        );
+        serviceConfiguration.updateConfig(disabled);
+        assertFalse(serviceConfiguration.isRealtimeUpdatesEnabled());
+    }
+
+    @Test
+    public void testIsRealtimeUpdatesEnabledFalseForUnrelatedFeature() {
+        APIResponse apiResponse = new APIResponse(
+                true, null, 0L, null, null, null,
+                false, null, true, false, null, List.of("some_other_feature")
+        );
+        serviceConfiguration.updateConfig(apiResponse);
+
+        assertFalse(serviceConfiguration.isRealtimeUpdatesEnabled());
+    }
 }