Add realtime config updates via SSE

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

@@ -1,11 +1,13 @@
 package dev.aikido.agent_api.background;
 
 import dev.aikido.agent_api.background.cloud.RealtimeAPI;
+import dev.aikido.agent_api.background.cloud.RealtimeSSEAPI;
 import dev.aikido.agent_api.background.cloud.api.APIResponse;
 import dev.aikido.agent_api.background.cloud.api.ReportingApi;
 import dev.aikido.agent_api.background.cloud.api.ReportingApiHTTP;
 import dev.aikido.agent_api.background.cloud.api.events.Started;
 import dev.aikido.agent_api.helpers.env.BlockingEnv;
+import dev.aikido.agent_api.helpers.env.FeatureFlags;
 import dev.aikido.agent_api.helpers.env.Token;
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
@@ -56,5 +58,9 @@ 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();
+        }
     }
 }

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

@@ -0,0 +1,63 @@
+package dev.aikido.agent_api.background;
+
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
+import dev.aikido.agent_api.background.cloud.RealtimeSSEAPI;
+import dev.aikido.agent_api.background.cloud.SSEParser;
+import dev.aikido.agent_api.background.cloud.api.APIResponse;
+import dev.aikido.agent_api.background.cloud.api.ReportingApi;
+import dev.aikido.agent_api.background.cloud.api.ReportingApiHTTP;
+import dev.aikido.agent_api.helpers.logging.LogManager;
+import dev.aikido.agent_api.helpers.logging.Logger;
+import dev.aikido.agent_api.storage.ServiceConfigStore;
+
+import java.util.Optional;
+
+public class RealtimeSSETask extends Thread {
+    private static final Logger logger = LogManager.getLogger(RealtimeSSETask.class);
+    private final RealtimeSSEAPI realtimeSSEApi;
+    private final ReportingApiHTTP reportingApi;
+    private Optional<Long> configLastUpdatedAt;
+
+    public RealtimeSSETask(RealtimeSSEAPI realtimeSSEApi, ReportingApiHTTP reportingApi) {
+        super("RealtimeSSETask");
+        this.realtimeSSEApi = realtimeSSEApi;
+        this.reportingApi = reportingApi;
+        this.configLastUpdatedAt = Optional.empty();
+        setDaemon(true);
+    }
+
+    @Override
+    public void run() {
+        realtimeSSEApi.listen(this::onEvent);
+    }
+
+    public void onEvent(SSEParser.Event event) {
+        logger.trace("SSE event received: %s", event.event());
+        if (!"config-updated".equals(event.event())) {
+            return;
+        }
+
+        long configUpdatedAt;
+        try {
+            JsonObject payload = JsonParser.parseString(event.data()).getAsJsonObject();
+            configUpdatedAt = payload.get("configUpdatedAt").getAsLong();
+        } catch (RuntimeException e) {
+            logger.debug("SSE config-updated event has invalid payload: %s", event.data());
+            return;
+        }
+
+        if (configLastUpdatedAt.isPresent() && configUpdatedAt <= configLastUpdatedAt.get()) {
+            return;
+        }
+        configLastUpdatedAt = Optional.of(configUpdatedAt);
+
+        Optional<APIResponse> newConfig = reportingApi.fetchNewConfig();
+        newConfig.ifPresent(ServiceConfigStore::updateFromAPIResponse);
+
+        Optional<ReportingApi.APIListsResponse> blockedListsRes = reportingApi.fetchBlockedLists();
+        blockedListsRes.ifPresent(ServiceConfigStore::updateFromAPIListsResponse);
+
+        logger.debug("Config updated via SSE");
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/RealtimeSSEAPI.java

@@ -0,0 +1,129 @@
+package dev.aikido.agent_api.background.cloud;
+
+import dev.aikido.agent_api.Config;
+import dev.aikido.agent_api.helpers.env.Token;
+import dev.aikido.agent_api.helpers.logging.LogManager;
+import dev.aikido.agent_api.helpers.logging.Logger;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.net.HttpURLConnection;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.Optional;
+import java.util.function.Consumer;
+
+import static dev.aikido.agent_api.helpers.env.Endpoints.getAikidoAPIEndpoint;
+
+public class RealtimeSSEAPI {
+    private static final Logger logger = LogManager.getLogger(RealtimeSSEAPI.class);
+    private static final int DEFAULT_INITIAL_RECONNECT_MS = 5_000;
+    private static final int DEFAULT_MAX_RECONNECT_MS = 60_000;
+    private static final int DEFAULT_STABLE_CONNECTION_MS = 30_000;
+    private static final int DEFAULT_READ_TIMEOUT_MS = 70_000;
+    private static final int CONNECT_TIMEOUT_MS = 10_000;
+
+    private final Token token;
+    private final String realtimeEndpoint;
+    private final int initialReconnectMs;
+    private final int maxReconnectMs;
+    private final int stableConnectionMs;
+    private final int readTimeoutMs;
+
+    public RealtimeSSEAPI(Token token) {
+        this(token, getAikidoAPIEndpoint(token), DEFAULT_INITIAL_RECONNECT_MS, DEFAULT_MAX_RECONNECT_MS, DEFAULT_STABLE_CONNECTION_MS, DEFAULT_READ_TIMEOUT_MS);
+    }
+
+    public RealtimeSSEAPI(Token token, String realtimeEndpoint, int initialReconnectMs, int maxReconnectMs, int stableConnectionMs, int readTimeoutMs) {
+        this.token = token;
+        this.realtimeEndpoint = realtimeEndpoint;
+        this.initialReconnectMs = initialReconnectMs;
+        this.maxReconnectMs = maxReconnectMs;
+        this.stableConnectionMs = stableConnectionMs;
+        this.readTimeoutMs = readTimeoutMs;
+    }
+
+    private enum Outcome { ERROR, DISCONNECTED, REJECTED }
+
+    private record ConnectResult(Outcome outcome, int statusCode) {}
+
+    public void listen(Consumer<SSEParser.Event> onEvent) {
+        int reconnectMs = initialReconnectMs;
+        while (!Thread.currentThread().isInterrupted()) {
+            long start = System.currentTimeMillis();
+            ConnectResult result = connect(onEvent);
+
+            if (result.outcome() == Outcome.REJECTED) {
+                logger.info("SSE connection rejected with status %s, stopping", result.statusCode());
+                return;
+            }
+
+            if (System.currentTimeMillis() - start >= stableConnectionMs) {
+                reconnectMs = initialReconnectMs;
+            }
+
+            long jitter = (long) (Math.random() * (reconnectMs / 2.0));
+            long delayMs = reconnectMs + jitter;
+            reconnectMs = Math.min(reconnectMs * 2, maxReconnectMs);
+
+            logger.trace("SSE scheduling reconnect in %sms", delayMs);
+            try {
+                Thread.sleep(delayMs);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                return;
+            }
+        }
+    }
+
+    private ConnectResult connect(Consumer<SSEParser.Event> onEvent) {
+        logger.trace("SSE connecting to realtime endpoint");
+        HttpURLConnection connection;
+        try {
+            URI uri = URI.create(realtimeEndpoint + "api/runtime/stream");
+            connection = (HttpURLConnection) uri.toURL().openConnection();
+            connection.setRequestMethod("GET");
+            connection.setRequestProperty("Authorization", token.get());
+            connection.setRequestProperty("Accept", "text/event-stream");
+            connection.setRequestProperty("Cache-Control", "no-cache");
+            connection.setRequestProperty("X-Agent-Platform", "java");
+            connection.setRequestProperty("X-Agent-Version", Config.pkgVersion);
+            connection.setConnectTimeout(CONNECT_TIMEOUT_MS);
+            connection.setReadTimeout(readTimeoutMs);
+        } catch (IOException e) {
+            logger.debug("SSE connection error: %s", e.getMessage());
+            return new ConnectResult(Outcome.ERROR, -1);
+        }
+
+        try {
+            int statusCode = connection.getResponseCode();
+            if (statusCode != 200) {
+                if (statusCode == 401 || statusCode == 403) {
+                    return new ConnectResult(Outcome.REJECTED, statusCode);
+                }
+                return new ConnectResult(Outcome.DISCONNECTED, statusCode);
+            }
+
+            logger.debug("SSE connected successfully");
+            try (BufferedReader reader = new BufferedReader(
+                    new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8))) {
+                SSEParser parser = new SSEParser(reader);
+                Optional<SSEParser.Event> event;
+                while ((event = parser.nextEvent()).isPresent()) {
+                    onEvent.accept(event.get());
+                }
+            }
+            logger.debug("SSE connection closed by server");
+            return new ConnectResult(Outcome.DISCONNECTED, statusCode);
+        } catch (IOException e) {
+            logger.debug("SSE stream error: %s", e.getMessage());
+            return new ConnectResult(Outcome.ERROR, -1);
+        } catch (RuntimeException e) {
+            logger.debug("SSE parser or callback error: %s", e.getMessage());
+            return new ConnectResult(Outcome.ERROR, -1);
+        } finally {
+            connection.disconnect();
+        }
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/SSEParser.java

@@ -0,0 +1,59 @@
+package dev.aikido.agent_api.background.cloud;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.util.Optional;
+
+public class SSEParser {
+    private final BufferedReader reader;
+
+    public SSEParser(BufferedReader reader) {
+        this.reader = reader;
+    }
+
+    public record Event(String event, String data) {}
+
+    public Optional<Event> nextEvent() throws IOException {
+        String eventType = "";
+        StringBuilder data = new StringBuilder();
+        boolean hasData = false;
+
+        String line;
+        while ((line = reader.readLine()) != null) {
+            if (line.isEmpty()) {
+                if (hasData) {
+                    if (data.charAt(data.length() - 1) == '\n') {
+                        data.setLength(data.length() - 1);
+                    }
+                    return Optional.of(new Event(eventType, data.toString()));
+                }
+                eventType = "";
+                data.setLength(0);
+                continue;
+            }
+            if (line.startsWith(":")) {
+                continue;
+            }
+
+            int sep = line.indexOf(':');
+            String field = sep == -1 ? line : line.substring(0, sep);
+            String value;
+            if (sep == -1) {
+                value = "";
+            } else {
+                value = line.substring(sep + 1);
+                if (value.startsWith(" ")) {
+                    value = value.substring(1);
+                }
+            }
+
+            if (field.equals("event")) {
+                eventType = value;
+            } else if (field.equals("data")) {
+                data.append(value).append("\n");
+                hasData = true;
+            }
+        }
+        return Optional.empty();
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/helpers/env/FeatureFlags.java

@@ -0,0 +1,9 @@
+package dev.aikido.agent_api.helpers.env;
+
+public enum FeatureFlags {
+    AIKIDO_FEATURE_SSE;
+
+    public boolean isEnabled() {
+        return new BooleanEnv(name(), false).getValue();
+    }
+}

agent_api/src/test/java/background/RealtimeSSETaskTest.java

@@ -0,0 +1,84 @@
+package background;
+
+import dev.aikido.agent_api.background.RealtimeSSETask;
+import dev.aikido.agent_api.background.cloud.RealtimeSSEAPI;
+import dev.aikido.agent_api.background.cloud.SSEParser;
+import dev.aikido.agent_api.background.cloud.api.APIResponse;
+import dev.aikido.agent_api.background.cloud.api.ReportingApi;
+import dev.aikido.agent_api.background.cloud.api.ReportingApiHTTP;
+import dev.aikido.agent_api.helpers.env.Token;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Optional;
+
+import static org.mockito.Mockito.*;
+
+public class RealtimeSSETaskTest {
+    private ReportingApiHTTP reportingApi;
+    private RealtimeSSETask task;
+
+    @BeforeEach
+    public void setUp() {
+        reportingApi = mock(ReportingApiHTTP.class);
+        task = new RealtimeSSETask(new RealtimeSSEAPI(new Token("token")), reportingApi);
+    }
+
+    private APIResponse sampleConfig(long configUpdatedAt) {
+        return new APIResponse(true, null, configUpdatedAt, List.of(), List.of(), List.of(), false, List.of(), true, true, List.of());
+    }
+
+    @Test
+    public void testIgnoresEventsThatAreNotConfigUpdated() {
+        task.onEvent(new SSEParser.Event("ping", ""));
+
+        verify(reportingApi, never()).fetchNewConfig();
+    }
+
+    @Test
+    public void testIgnoresInvalidJsonPayload() {
+        task.onEvent(new SSEParser.Event("config-updated", "not json"));
+
+        verify(reportingApi, never()).fetchNewConfig();
+    }
+
+    @Test
+    public void testIgnoresPayloadMissingConfigUpdatedAt() {
+        task.onEvent(new SSEParser.Event("config-updated", "{\"foo\":\"bar\"}"));
+
+        verify(reportingApi, never()).fetchNewConfig();
+    }
+
+    @Test
+    public void testFetchesAndAppliesNewConfigOnNewerEvent() {
+        when(reportingApi.fetchNewConfig()).thenReturn(Optional.of(sampleConfig(200)));
+        when(reportingApi.fetchBlockedLists()).thenReturn(Optional.empty());
+
+        task.onEvent(new SSEParser.Event("config-updated", "{\"configUpdatedAt\":200}"));
+
+        verify(reportingApi, times(1)).fetchNewConfig();
+        verify(reportingApi, times(1)).fetchBlockedLists();
+    }
+
+    @Test
+    public void testIgnoresEventThatIsNotNewer() {
+        when(reportingApi.fetchNewConfig()).thenReturn(Optional.of(sampleConfig(200)));
+        when(reportingApi.fetchBlockedLists()).thenReturn(Optional.empty());
+
+        task.onEvent(new SSEParser.Event("config-updated", "{\"configUpdatedAt\":200}"));
+        task.onEvent(new SSEParser.Event("config-updated", "{\"configUpdatedAt\":150}"));
+
+        verify(reportingApi, times(1)).fetchNewConfig();
+    }
+
+    @Test
+    public void testHandlesFetchFailureGracefully() {
+        when(reportingApi.fetchNewConfig()).thenReturn(Optional.empty());
+        when(reportingApi.fetchBlockedLists()).thenReturn(Optional.empty());
+
+        task.onEvent(new SSEParser.Event("config-updated", "{\"configUpdatedAt\":200}"));
+
+        verify(reportingApi, times(1)).fetchNewConfig();
+    }
+}