From 1654765b05409e52436107219d431a8dfb7387da Mon Sep 17 00:00:00 2001 From: dmitrii Date: Wed, 22 Jul 2026 18:31:13 +0200 Subject: [PATCH] Add realtime config updates via SSE Agent subscribes to a Server-Sent Events stream and applies config changes as soon as the cloud pushes a config-updated event, instead of waiting for the next 60s poll. The existing RealtimeTask poll stays untouched and keeps running as a fallback. Gated behind AIKIDO_FEATURE_SSE (off by default), matching the firewall-node and firewall-python ports of this feature. Uses HttpURLConnection for the stream connection since its read timeout is a true idle timeout, unlike java.net.http.HttpClient's whole-exchange timeout. --- .../background/BackgroundProcess.java | 6 + .../agent_api/background/RealtimeSSETask.java | 63 ++++ .../background/cloud/RealtimeSSEAPI.java | 129 +++++++ .../agent_api/background/cloud/SSEParser.java | 59 ++++ .../agent_api/helpers/env/FeatureFlags.java | 9 + .../java/background/RealtimeSSETaskTest.java | 84 +++++ .../background/cloud/RealtimeSSEAPITest.java | 323 ++++++++++++++++++ .../java/background/cloud/SSEParserTest.java | 108 ++++++ 8 files changed, 781 insertions(+) create mode 100644 agent_api/src/main/java/dev/aikido/agent_api/background/RealtimeSSETask.java create mode 100644 agent_api/src/main/java/dev/aikido/agent_api/background/cloud/RealtimeSSEAPI.java create mode 100644 agent_api/src/main/java/dev/aikido/agent_api/background/cloud/SSEParser.java create mode 100644 agent_api/src/main/java/dev/aikido/agent_api/helpers/env/FeatureFlags.java create mode 100644 agent_api/src/test/java/background/RealtimeSSETaskTest.java create mode 100644 agent_api/src/test/java/background/cloud/RealtimeSSEAPITest.java create mode 100644 agent_api/src/test/java/background/cloud/SSEParserTest.java diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java b/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java index 02419abe..5baaa1f6 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java +++ b/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(); + } } } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/RealtimeSSETask.java b/agent_api/src/main/java/dev/aikido/agent_api/background/RealtimeSSETask.java new file mode 100644 index 00000000..925cc775 --- /dev/null +++ b/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 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 newConfig = reportingApi.fetchNewConfig(); + newConfig.ifPresent(ServiceConfigStore::updateFromAPIResponse); + + Optional blockedListsRes = reportingApi.fetchBlockedLists(); + blockedListsRes.ifPresent(ServiceConfigStore::updateFromAPIListsResponse); + + logger.debug("Config updated via SSE"); + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/RealtimeSSEAPI.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/RealtimeSSEAPI.java new file mode 100644 index 00000000..cad6d188 --- /dev/null +++ b/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 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 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 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(); + } + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/SSEParser.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/SSEParser.java new file mode 100644 index 00000000..128c731c --- /dev/null +++ b/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 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(); + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/FeatureFlags.java b/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/FeatureFlags.java new file mode 100644 index 00000000..dcf42613 --- /dev/null +++ b/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(); + } +} diff --git a/agent_api/src/test/java/background/RealtimeSSETaskTest.java b/agent_api/src/test/java/background/RealtimeSSETaskTest.java new file mode 100644 index 00000000..8b304061 --- /dev/null +++ b/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(); + } +} diff --git a/agent_api/src/test/java/background/cloud/RealtimeSSEAPITest.java b/agent_api/src/test/java/background/cloud/RealtimeSSEAPITest.java new file mode 100644 index 00000000..ca44d2cd --- /dev/null +++ b/agent_api/src/test/java/background/cloud/RealtimeSSEAPITest.java @@ -0,0 +1,323 @@ +package background.cloud; + +import com.sun.net.httpserver.HttpServer; +import dev.aikido.agent_api.background.cloud.RealtimeSSEAPI; +import dev.aikido.agent_api.background.cloud.SSEParser; +import dev.aikido.agent_api.helpers.env.Token; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.net.ServerSocket; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.*; + +public class RealtimeSSEAPITest { + private HttpServer server; + + @AfterEach + public void tearDown() { + if (server != null) { + server.stop(0); + } + } + + private int freePort() throws IOException { + try (ServerSocket socket = new ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private String endpointFor(int port) { + return "http://localhost:" + port + "/"; + } + + private void writeEvent(OutputStream os, String event, String data) throws IOException { + String payload = "event: " + event + "\ndata: " + data + "\n\n"; + os.write(payload.getBytes(StandardCharsets.UTF_8)); + os.flush(); + } + + @Test + public void testReceivesConfigUpdatedEvent() throws Exception { + int port = freePort(); + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + exchange.getResponseHeaders().add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + writeEvent(exchange.getResponseBody(), "config-updated", "{\"configUpdatedAt\":1}"); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 50, 200, 1000, 5000); + + BlockingQueue events = new ArrayBlockingQueue<>(10); + Thread listener = new Thread(() -> api.listen(events::add)); + listener.setDaemon(true); + listener.start(); + + SSEParser.Event event = events.poll(5, TimeUnit.SECONDS); + listener.interrupt(); + listener.join(2000); + + assertNotNull(event); + assertEquals("config-updated", event.event()); + assertEquals("{\"configUpdatedAt\":1}", event.data()); + } + + @Test + public void testStopsRetryingOn401() throws Exception { + int port = freePort(); + AtomicInteger hits = new AtomicInteger(0); + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + hits.incrementAndGet(); + exchange.sendResponseHeaders(401, -1); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 50, 200, 1000, 5000); + + Thread listener = new Thread(() -> api.listen(event -> {})); + listener.setDaemon(true); + listener.start(); + listener.join(2000); + + assertFalse(listener.isAlive()); + assertEquals(1, hits.get()); + } + + @Test + public void testReconnectsAfter500() throws Exception { + int port = freePort(); + AtomicInteger hits = new AtomicInteger(0); + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + if (hits.getAndIncrement() == 0) { + exchange.sendResponseHeaders(500, -1); + exchange.close(); + return; + } + exchange.getResponseHeaders().add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + writeEvent(exchange.getResponseBody(), "config-updated", "{\"configUpdatedAt\":2}"); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 50, 200, 1000, 5000); + + BlockingQueue events = new ArrayBlockingQueue<>(10); + Thread listener = new Thread(() -> api.listen(events::add)); + listener.setDaemon(true); + listener.start(); + + SSEParser.Event event = events.poll(5, TimeUnit.SECONDS); + listener.interrupt(); + listener.join(2000); + + assertNotNull(event); + assertEquals(2, hits.get()); + } + + @Test + public void testReconnectsAfterReadTimeout() throws Exception { + int port = freePort(); + AtomicInteger hits = new AtomicInteger(0); + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + if (hits.getAndIncrement() == 0) { + exchange.getResponseHeaders().add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + return; + } + exchange.getResponseHeaders().add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + writeEvent(exchange.getResponseBody(), "config-updated", "{\"configUpdatedAt\":3}"); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 50, 200, 1000, 300); + + BlockingQueue events = new ArrayBlockingQueue<>(10); + Thread listener = new Thread(() -> api.listen(events::add)); + listener.setDaemon(true); + listener.start(); + + SSEParser.Event event = events.poll(5, TimeUnit.SECONDS); + listener.interrupt(); + listener.join(2000); + + assertNotNull(event); + assertEquals(2, hits.get()); + } + + @Test + public void testConnectionRefusedRetries() throws Exception { + int port = freePort(); + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 50, 200, 1000, 5000); + + Thread listener = new Thread(() -> api.listen(event -> {})); + listener.setDaemon(true); + listener.start(); + Thread.sleep(300); + listener.interrupt(); + listener.join(2000); + + assertFalse(listener.isAlive()); + } + + @Test + public void testSendsExpectedRequestHeaders() throws Exception { + int port = freePort(); + List receivedHeaders = new CopyOnWriteArrayList<>(); + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + receivedHeaders.add(exchange.getRequestHeaders()); + exchange.sendResponseHeaders(401, -1); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("my-token"), endpointFor(port), 50, 200, 1000, 5000); + Thread listener = new Thread(() -> api.listen(event -> {})); + listener.setDaemon(true); + listener.start(); + listener.join(2000); + + assertEquals(1, receivedHeaders.size()); + com.sun.net.httpserver.Headers headers = receivedHeaders.get(0); + assertEquals("my-token", headers.getFirst("Authorization")); + assertEquals("text/event-stream", headers.getFirst("Accept")); + assertEquals("no-cache", headers.getFirst("Cache-Control")); + assertEquals("java", headers.getFirst("X-Agent-Platform")); + assertNotNull(headers.getFirst("X-Agent-Version")); + } + + @Test + public void testCallbackExceptionIsTreatedAsRetryableError() throws Exception { + int port = freePort(); + AtomicInteger hits = new AtomicInteger(0); + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + hits.incrementAndGet(); + exchange.getResponseHeaders().add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + writeEvent(exchange.getResponseBody(), "config-updated", "{\"configUpdatedAt\":1}"); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 20, 100, 1000, 5000); + Thread listener = new Thread(() -> api.listen(event -> { + throw new RuntimeException("callback exploded"); + })); + listener.setDaemon(true); + listener.start(); + + long deadline = System.currentTimeMillis() + 2000; + while (hits.get() < 3 && System.currentTimeMillis() < deadline) { + Thread.sleep(20); + } + listener.interrupt(); + listener.join(2000); + + assertTrue(hits.get() >= 3, "expected multiple reconnects after callback exceptions, got " + hits.get()); + } + + @Test + public void testBackoffDoublesBetweenReconnects() throws Exception { + int port = freePort(); + List hitTimestamps = new CopyOnWriteArrayList<>(); + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + hitTimestamps.add(System.currentTimeMillis()); + exchange.sendResponseHeaders(500, -1); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 150, 5000, 100_000, 5000); + Thread listener = new Thread(() -> api.listen(event -> {})); + listener.setDaemon(true); + listener.start(); + + long deadline = System.currentTimeMillis() + 5000; + while (hitTimestamps.size() < 4 && System.currentTimeMillis() < deadline) { + Thread.sleep(20); + } + listener.interrupt(); + listener.join(2000); + + assertTrue(hitTimestamps.size() >= 4, "expected at least 4 attempts, got " + hitTimestamps.size()); + long gap1 = hitTimestamps.get(1) - hitTimestamps.get(0); + long gap2 = hitTimestamps.get(2) - hitTimestamps.get(1); + long gap3 = hitTimestamps.get(3) - hitTimestamps.get(2); + + assertTrue(gap2 > gap1 * 1.2, "expected gap2 (" + gap2 + "ms) to be noticeably larger than gap1 (" + gap1 + "ms)"); + assertTrue(gap3 > gap2 * 1.2, "expected gap3 (" + gap3 + "ms) to be noticeably larger than gap2 (" + gap2 + "ms)"); + } + + @Test + public void testBackoffResetsAfterStableConnection() throws Exception { + int port = freePort(); + List hitTimestamps = new CopyOnWriteArrayList<>(); + AtomicInteger hits = new AtomicInteger(0); + int holdMs = 50; + server = HttpServer.create(new InetSocketAddress(port), 0); + server.createContext("/api/runtime/stream", exchange -> { + hitTimestamps.add(System.currentTimeMillis()); + int hitNumber = hits.getAndIncrement(); + if (hitNumber < 3) { + exchange.sendResponseHeaders(500, -1); + exchange.close(); + return; + } + if (hitNumber == 3) { + exchange.getResponseHeaders().add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + try { + Thread.sleep(holdMs); + } catch (InterruptedException ignored) { + } + exchange.close(); + return; + } + exchange.sendResponseHeaders(500, -1); + exchange.close(); + }); + server.start(); + + RealtimeSSEAPI api = new RealtimeSSEAPI(new Token("token"), endpointFor(port), 80, 5000, 30, 5000); + Thread listener = new Thread(() -> api.listen(event -> {})); + listener.setDaemon(true); + listener.start(); + + long deadline = System.currentTimeMillis() + 5000; + while (hitTimestamps.size() < 5 && System.currentTimeMillis() < deadline) { + Thread.sleep(20); + } + listener.interrupt(); + listener.join(2000); + + assertTrue(hitTimestamps.size() >= 5, "expected at least 5 attempts, got " + hitTimestamps.size()); + long gapBeforeReset = hitTimestamps.get(3) - hitTimestamps.get(2); + long gapAfterReset = hitTimestamps.get(4) - hitTimestamps.get(3); + + assertTrue(gapAfterReset < gapBeforeReset, + "expected backoff to reset after a stable connection instead of keep growing: gapBeforeReset=" + + gapBeforeReset + "ms gapAfterReset=" + gapAfterReset + "ms"); + } +} diff --git a/agent_api/src/test/java/background/cloud/SSEParserTest.java b/agent_api/src/test/java/background/cloud/SSEParserTest.java new file mode 100644 index 00000000..002418ce --- /dev/null +++ b/agent_api/src/test/java/background/cloud/SSEParserTest.java @@ -0,0 +1,108 @@ +package background.cloud; + +import dev.aikido.agent_api.background.cloud.SSEParser; +import org.junit.jupiter.api.Test; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.StringReader; +import java.util.Optional; + +import static org.junit.jupiter.api.Assertions.*; + +public class SSEParserTest { + private SSEParser parserFor(String raw) { + return new SSEParser(new BufferedReader(new StringReader(raw))); + } + + @Test + public void testParsesEventAndData() throws IOException { + SSEParser parser = parserFor("event: config-updated\ndata: {\"configUpdatedAt\":1}\n\n"); + + Optional event = parser.nextEvent(); + + assertTrue(event.isPresent()); + assertEquals("config-updated", event.get().event()); + assertEquals("{\"configUpdatedAt\":1}", event.get().data()); + } + + @Test + public void testJoinsMultiLineData() throws IOException { + SSEParser parser = parserFor("event: config-updated\ndata: line1\ndata: line2\n\n"); + + Optional event = parser.nextEvent(); + + assertTrue(event.isPresent()); + assertEquals("line1\nline2", event.get().data()); + } + + @Test + public void testIgnoresCommentLines() throws IOException { + SSEParser parser = parserFor(": ping\nevent: config-updated\ndata: hello\n\n"); + + Optional event = parser.nextEvent(); + + assertTrue(event.isPresent()); + assertEquals("hello", event.get().data()); + } + + @Test + public void testPureCommentDoesNotDispatch() throws IOException { + SSEParser parser = parserFor(": ping\n\nevent: config-updated\ndata: hello\n\n"); + + Optional event = parser.nextEvent(); + + assertTrue(event.isPresent()); + assertEquals("hello", event.get().data()); + } + + @Test + public void testIncompleteEventAtStreamEndIsDiscarded() throws IOException { + SSEParser parser = parserFor("event: config-updated\ndata: hello"); + + Optional event = parser.nextEvent(); + + assertTrue(event.isEmpty()); + } + + @Test + public void testMultipleEventsInSequence() throws IOException { + SSEParser parser = parserFor("data: first\n\ndata: second\n\n"); + + Optional first = parser.nextEvent(); + Optional second = parser.nextEvent(); + Optional third = parser.nextEvent(); + + assertEquals("first", first.get().data()); + assertEquals("second", second.get().data()); + assertTrue(third.isEmpty()); + } + + @Test + public void testValueWithoutLeadingSpaceIsPreserved() throws IOException { + SSEParser parser = parserFor("data:hello\n\n"); + + Optional event = parser.nextEvent(); + + assertEquals("hello", event.get().data()); + } + + @Test + public void testUnknownFieldIsIgnored() throws IOException { + SSEParser parser = parserFor("id: 42\ndata: hello\n\n"); + + Optional event = parser.nextEvent(); + + assertTrue(event.isPresent()); + assertEquals("hello", event.get().data()); + } + + @Test + public void testEmptyStreamReturnsEmpty() throws IOException { + SSEParser parser = parserFor(""); + + Optional event = parser.nextEvent(); + + assertTrue(event.isEmpty()); + } +}