Skip to content

Support gradual rollout flag for realtime updates #328

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mishenevd wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Support gradual rollout flag for realtime updates #328

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ 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()) {
if (token != null && (FeatureFlags.AIKIDO_FEATURE_SSE.isEnabled() || ServiceConfigStore.isRealtimeUpdatesEnabled())) {
new RealtimeSSETask(new RealtimeSSEAPI(token), api).start();
}
}
Expand Down
19 changes: 18 additions & 1 deletion agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,23 @@ public record APIResponse(
List<Domain> domains,
boolean receivedAnyStats,
boolean block,
List<String> excludedUserIdsFromRateLimiting
List<String> excludedUserIdsFromRateLimiting,
List<String> enabledFeatures

@timokoessler timokoessler Jul 31, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might be possible that this becomes null if not set in Core response? And then new HashSet<>(null) might throw? Not a Java export tho.

) {
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());
}
}
9 changes: 9 additions & 0 deletions agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfigStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
8 changes: 8 additions & 0 deletions agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand All @@ -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();
}
Expand All @@ -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;
}
Expand Down
14 changes: 14 additions & 0 deletions agent_api/src/test/java/storage/ServiceConfigStoreTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}
44 changes: 44 additions & 0 deletions agent_api/src/test/java/storage/ServiceConfigurationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -665,4 +665,48 @@ public void testIsIpBlockedWithAllowedIPsAndBlockedIPs() {
assertTrue(resultBlocked.blocked());
assertFalse(resultNotAllowedLocal.blocked());
}

@Test
public void testIsRealtimeUpdatesEnabledDefaultsFalse() {
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());
}
}
Loading