From 44fc973011d8875b64ffe66bc6cc9b73828d5660 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Thu, 16 Jul 2026 12:00:21 +0200 Subject: [PATCH 1/4] Add config updates via SSE --- .../cloud_connection_manager/__init__.py | 7 +- .../background_process/realtime/__init__.py | 2 +- .../realtime/listen_for_config_updates.py | 59 +++++ .../listen_for_config_updates_test.py | 169 ++++++++++++++ .../realtime/sse_client/__init__.py | 127 ++++++++++ .../realtime/sse_client/init_test.py | 221 ++++++++++++++++++ .../realtime/sse_client/parser.py | 127 ++++++++++ .../realtime/sse_client/parser_test.py | 186 +++++++++++++++ 8 files changed, 896 insertions(+), 2 deletions(-) create mode 100644 aikido_zen/background_process/realtime/listen_for_config_updates.py create mode 100644 aikido_zen/background_process/realtime/listen_for_config_updates_test.py create mode 100644 aikido_zen/background_process/realtime/sse_client/__init__.py create mode 100644 aikido_zen/background_process/realtime/sse_client/init_test.py create mode 100644 aikido_zen/background_process/realtime/sse_client/parser.py create mode 100644 aikido_zen/background_process/realtime/sse_client/parser_test.py diff --git a/aikido_zen/background_process/cloud_connection_manager/__init__.py b/aikido_zen/background_process/cloud_connection_manager/__init__.py index ea92f5074..4e69037c4 100644 --- a/aikido_zen/background_process/cloud_connection_manager/__init__.py +++ b/aikido_zen/background_process/cloud_connection_manager/__init__.py @@ -1,4 +1,4 @@ -""" This file simply exports the CloudConnectionManager class""" +"""This file simply exports the CloudConnectionManager class""" from aikido_zen.background_process.heartbeats import send_heartbeats_every_x_secs from aikido_zen.background_process.routes import Routes @@ -10,10 +10,12 @@ from aikido_zen.storage.users import Users from aikido_zen.storage.hostnames import Hostnames from ..realtime.start_polling_for_changes import start_polling_for_changes +from ..realtime.listen_for_config_updates import listen_for_config_updates from ...helpers.get_current_unixtime_ms import get_unixtime_ms from ...storage.ai_statistics import AIStatistics from ...storage.firewall_lists import FirewallLists from ...storage.statistics import Statistics +from aikido_zen.helpers.env_vars.feature_flags import is_feature_enabled # Import functions : from .get_manager_info import get_manager_info @@ -67,6 +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"): + listen_for_config_updates(self, event_scheduler) + def report_initial_stats(self): """ This is run 1m after startup, and checks if we should send out diff --git a/aikido_zen/background_process/realtime/__init__.py b/aikido_zen/background_process/realtime/__init__.py index e50e08106..5f12f58bf 100644 --- a/aikido_zen/background_process/realtime/__init__.py +++ b/aikido_zen/background_process/realtime/__init__.py @@ -15,7 +15,7 @@ def get_realtime_url(): if not realtime_url.endswith("/"): realtime_url += "/" return realtime_url - return "https://runtime.aikido.dev/" + return get_api_url() def get_config(token): diff --git a/aikido_zen/background_process/realtime/listen_for_config_updates.py b/aikido_zen/background_process/realtime/listen_for_config_updates.py new file mode 100644 index 000000000..5253a4d18 --- /dev/null +++ b/aikido_zen/background_process/realtime/listen_for_config_updates.py @@ -0,0 +1,59 @@ +""" +Mainly exports `listen_for_config_updates` +""" + +import json + +from aikido_zen.helpers.token import Token +from aikido_zen.helpers.logging import logger +import aikido_zen.background_process.realtime as realtime +from .sse_client import connect_to_sse + + +def listen_for_config_updates(connection_manager, event_scheduler): + """ + Connects to the realtime SSE endpoint and fetches the new config whenever + the server signals, through a "config-updated" event, that it has changed. + """ + if not isinstance(connection_manager.token, Token): + logger.info("No token provided, not listening for config updates") + return + if connection_manager.serverless: + logger.info( + "Running in serverless environment, not listening for config updates" + ) + return + + token = connection_manager.token + last_updated_at = connection_manager.conf.last_updated_at + + def on_event(event): + nonlocal last_updated_at + + logger.debug("SSE event received: %s", event.event) + if event.event != "config-updated": + return + + try: + payload = json.loads(event.data) + config_updated_at = payload["configUpdatedAt"] + if config_updated_at <= last_updated_at: + return + except (ValueError, KeyError, TypeError): + logger.debug("SSE config-updated event has invalid payload: %s", event.data) + return + + logger.debug("SSE config-updated event, fetching new config") + + try: + config = realtime.get_config(token) + logger.debug( + "SSE config fetched, configUpdatedAt: %s", config.get("configUpdatedAt") + ) + last_updated_at = config.get("configUpdatedAt", config_updated_at) + connection_manager.update_service_config({**config, "success": True}) + connection_manager.update_firewall_lists() + except Exception as e: + logger.error("Failed to fetch config after SSE event : %s", e) + + connect_to_sse(token=token, on_event=on_event) diff --git a/aikido_zen/background_process/realtime/listen_for_config_updates_test.py b/aikido_zen/background_process/realtime/listen_for_config_updates_test.py new file mode 100644 index 000000000..575f6d7c6 --- /dev/null +++ b/aikido_zen/background_process/realtime/listen_for_config_updates_test.py @@ -0,0 +1,169 @@ +import json +import logging +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from aikido_zen.helpers.token import Token +from .listen_for_config_updates import listen_for_config_updates + + +def make_event(event="config-updated", data=None): + return SimpleNamespace( + event=event, data=json.dumps(data) if data is not None else "" + ) + + +@pytest.fixture +def connection_manager(): + return MagicMock( + token=Token("123"), + serverless=None, + conf=MagicMock(last_updated_at=0), + ) + + +def test_no_token(caplog): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=MagicMock(token=None, serverless=None), + event_scheduler=MagicMock(), + ) + + assert "No token provided, not listening for config updates" in caplog.text + mock_connect.assert_not_called() + + +def test_serverless_environment(caplog): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=MagicMock(token=Token("123"), serverless=True), + event_scheduler=MagicMock(), + ) + + assert ( + "Running in serverless environment, not listening for config updates" + in caplog.text + ) + mock_connect.assert_not_called() + + +def test_connects_to_sse_with_token(connection_manager): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=connection_manager, event_scheduler=MagicMock() + ) + + mock_connect.assert_called_once() + assert mock_connect.call_args.kwargs["token"] is connection_manager.token + assert callable(mock_connect.call_args.kwargs["on_event"]) + + +def get_on_event(connection_manager): + with patch( + "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse" + ) as mock_connect: + listen_for_config_updates( + connection_manager=connection_manager, event_scheduler=MagicMock() + ) + return mock_connect.call_args.kwargs["on_event"] + + +def test_ignores_events_that_are_not_config_updated(connection_manager): + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(make_event(event="ping")) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + + +def test_ignores_config_updated_event_with_invalid_json(connection_manager, caplog): + caplog.set_level(logging.DEBUG, logger="Zen") + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(SimpleNamespace(event="config-updated", data="not json")) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + assert "SSE config-updated event has invalid payload" in caplog.text + + +def test_ignores_config_updated_event_missing_config_updated_at(connection_manager): + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(make_event(data={"foo": "bar"})) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + + +def test_ignores_config_updated_event_that_is_not_newer(connection_manager): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config: + on_event(make_event(data={"configUpdatedAt": 100})) + + mock_get_config.assert_not_called() + connection_manager.update_service_config.assert_not_called() + + +def test_fetches_and_applies_new_config_on_newer_event(connection_manager): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + new_config = {"endpoints": [], "configUpdatedAt": 200} + with patch( + "aikido_zen.background_process.realtime.get_config", return_value=new_config + ) as mock_get_config: + on_event(make_event(data={"configUpdatedAt": 200})) + + mock_get_config.assert_called_once_with(connection_manager.token) + connection_manager.update_service_config.assert_called_once_with( + {**new_config, "success": True} + ) + connection_manager.update_firewall_lists.assert_called_once() + + +def test_updates_last_updated_at_so_a_second_stale_event_is_ignored( + connection_manager, +): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + new_config = {"endpoints": [], "configUpdatedAt": 200} + with patch( + "aikido_zen.background_process.realtime.get_config", return_value=new_config + ) as mock_get_config: + on_event(make_event(data={"configUpdatedAt": 200})) + # A second event claiming to be newer than the original lastUpdatedAt + # (100), but not newer than what we just fetched (200), should now be + # ignored without fetching again. + on_event(make_event(data={"configUpdatedAt": 150})) + + mock_get_config.assert_called_once() + + +def test_handles_get_config_failure_gracefully(connection_manager, caplog): + connection_manager.conf.last_updated_at = 100 + on_event = get_on_event(connection_manager) + + with patch( + "aikido_zen.background_process.realtime.get_config", + side_effect=ValueError("Request timed out"), + ): + on_event(make_event(data={"configUpdatedAt": 200})) + + connection_manager.update_service_config.assert_not_called() + assert "Failed to fetch config after SSE event" in caplog.text diff --git a/aikido_zen/background_process/realtime/sse_client/__init__.py b/aikido_zen/background_process/realtime/sse_client/__init__.py new file mode 100644 index 000000000..aa0ee10aa --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/__init__.py @@ -0,0 +1,127 @@ +""" +SSE (Server-Sent Events) client that connects to the realtime endpoint and +dispatches events, automatically reconnecting with exponential backoff and +jitter on disconnects. +""" + +import http.client +import random +import socket +import threading +import time +import urllib.error +import urllib.request + +import aikido_zen.background_process.realtime as realtime +from aikido_zen.helpers.logging import logger +from .parser import SSEParser + +INITIAL_RECONNECT_SECS = 5 +MAX_RECONNECT_SECS = 60 +STABLE_CONNECTION_SECS = 30 +READ_TIMEOUT_SECS = 70 + +_CONNECTION_ERRORS = (urllib.error.URLError, OSError, http.client.HTTPException) + + +def connect_to_sse( + token, + on_event, + initial_reconnect_secs=INITIAL_RECONNECT_SECS, + read_timeout_secs=READ_TIMEOUT_SECS, +): + """ + Starts a daemon thread that connects to the realtime SSE endpoint and calls + `on_event(event)` for every event received. Reconnects with exponential + backoff (and jitter) on disconnect, until the server rejects the connection + with a 401 or 403 status, at which point it stops. + + Returns the thread that was started. + """ + thread = threading.Thread( + target=_reconnect_loop, + args=(token, on_event, initial_reconnect_secs, read_timeout_secs), + daemon=True, + ) + thread.start() + return thread + + +def _reconnect_loop(token, on_event, initial_reconnect_secs, read_timeout_secs): + reconnect_secs = initial_reconnect_secs + try: + while True: + start = time.monotonic() + outcome, status_code = _connect(token, on_event, read_timeout_secs) + + if outcome == "disconnected" and status_code in (401, 403): + logger.info( + "SSE connection rejected with status %s, stopping", status_code + ) + return + + if time.monotonic() - start >= STABLE_CONNECTION_SECS: + reconnect_secs = initial_reconnect_secs + + jitter = random.random() * (reconnect_secs / 2) + delay_secs = reconnect_secs + jitter + + logger.debug("SSE scheduling reconnect in %sms", round(delay_secs * 1000)) + + reconnect_secs = min(reconnect_secs * 2, MAX_RECONNECT_SECS) + + time.sleep(delay_secs) + except Exception as e: + logger.error("SSE loop error : %s", e) + + +def _connect(token, on_event, read_timeout_secs): + """ + Opens a single SSE connection and dispatches events until the connection is + closed, errors out, or goes idle for longer than `read_timeout_secs`. + + Returns a tuple of (outcome, status_code), where outcome is either "error" + (the connection could not be made, or timed out) or "disconnected" (a + response was received, but the connection has since ended). + """ + url = f"{realtime.get_realtime_url()}api/runtime/stream" + logger.debug("SSE connecting to %s", url) + + request = urllib.request.Request( + url, + method="GET", + headers={ + "Authorization": str(token), + "Accept": "text/event-stream", + "Cache-Control": "no-cache", + }, + ) + + try: + with urllib.request.urlopen(request, timeout=read_timeout_secs) as response: + status_code = response.getcode() + if status_code != 200: + return "disconnected", status_code + + logger.debug("SSE connected successfully") + + parser = SSEParser(response) + try: + for event in parser.events(): + logger.debug("SSE received event : %s", event) + on_event(event) + except socket.timeout: + logger.debug("SSE read timeout") + return "error", None + except _CONNECTION_ERRORS as e: + logger.debug("SSE stream error : %s", e) + return "disconnected", status_code + except urllib.error.HTTPError as e: + logger.debug("SSE connection rejected with status %s", e.code) + return "disconnected", e.code + except _CONNECTION_ERRORS as e: + logger.debug("SSE connection error : %s", e) + return "error", None + + logger.debug("SSE connection closed by server") + return "disconnected", status_code diff --git a/aikido_zen/background_process/realtime/sse_client/init_test.py b/aikido_zen/background_process/realtime/sse_client/init_test.py new file mode 100644 index 000000000..9e72954c6 --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/init_test.py @@ -0,0 +1,221 @@ +import socket +import urllib.error +from unittest.mock import MagicMock, patch + +from aikido_zen.helpers.token import Token +from . import ( + _connect, + _reconnect_loop, + connect_to_sse, + INITIAL_RECONNECT_SECS, + READ_TIMEOUT_SECS, +) + +REALTIME_URL_PATCH = patch( + "aikido_zen.background_process.realtime.get_realtime_url", + return_value="https://runtime.aikido.dev/", +) + + +def make_response(status_code, chunks=None): + response = MagicMock() + response.getcode.return_value = status_code + response.__enter__.return_value = response + response.__exit__.side_effect = lambda *args: response.close() + if chunks is not None: + response.__iter__.return_value = iter(chunks) + return response + + +def test_connect_dispatches_events_and_reports_disconnected_on_stream_end(): + response = make_response(200, [b"data: hello\n\n", b"data: world\n\n"]) + events = [] + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), events.append, 5) + + assert outcome == "disconnected" + assert status_code == 200 + assert [e.data for e in events] == ["hello", "world"] + response.close.assert_called_once() # closed via parser.close() + + +def test_connect_sends_expected_request(): + response = make_response(200, []) + captured_request = {} + + def fake_urlopen(request, timeout): + captured_request["request"] = request + captured_request["timeout"] = timeout + return response + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", side_effect=fake_urlopen): + _connect(Token("my-token"), MagicMock(), 42) + + request = captured_request["request"] + assert request.full_url == "https://runtime.aikido.dev/api/runtime/stream" + assert request.get_header("Authorization") == "my-token" + assert request.get_header("Accept") == "text/event-stream" + assert captured_request["timeout"] == 42 + + +def test_connect_non_200_status_closes_response_without_reading_body(): + response = make_response(500) + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "disconnected" + assert status_code == 500 + response.close.assert_called_once() + + +def test_connect_http_error_returns_disconnected_with_status(): + http_error = urllib.error.HTTPError( + url="https://runtime.aikido.dev/api/runtime/stream", + code=403, + msg="Forbidden", + hdrs=None, + fp=None, + ) + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", side_effect=http_error): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "disconnected" + assert status_code == 403 + + +def test_connect_connection_error_returns_error_outcome(): + with REALTIME_URL_PATCH, patch( + "urllib.request.urlopen", + side_effect=urllib.error.URLError("connection refused"), + ): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "error" + assert status_code is None + + +def test_connect_read_timeout_mid_stream_returns_error(): + response = make_response(200) + response.__iter__.side_effect = socket.timeout("timed out") + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "error" + assert status_code is None + response.close.assert_called_once() + + +def test_connect_generic_stream_error_returns_disconnected_with_status(): + response = make_response(200) + response.__iter__.side_effect = ConnectionResetError("reset") + + with REALTIME_URL_PATCH, patch("urllib.request.urlopen", return_value=response): + outcome, status_code = _connect(Token("123"), MagicMock(), 5) + + assert outcome == "disconnected" + assert status_code == 200 + + +def test_reconnect_loop_stops_immediately_on_401_or_403(): + for status_code in (401, 403): + connect_calls = [] + + def fake_connect(token, on_event, read_timeout_secs): + connect_calls.append(status_code) + return "disconnected", status_code + + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=fake_connect, + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.sleep" + ) as mock_sleep: + _reconnect_loop(Token("123"), MagicMock(), INITIAL_RECONNECT_SECS, 5) + + assert connect_calls == [status_code] + mock_sleep.assert_not_called() + + +def test_reconnect_loop_backoff_doubles_between_reconnects(): + outcomes = iter( + [("error", None), ("error", None), ("disconnected", 401)], + ) + sleep_calls = [] + monotonic_values = iter([0, 0, 10, 10, 20]) + + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=lambda *a: next(outcomes), + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.sleep", + side_effect=sleep_calls.append, + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.monotonic", + side_effect=lambda: next(monotonic_values), + ), patch( + "aikido_zen.background_process.realtime.sse_client.random.random", + return_value=0, + ): + _reconnect_loop(Token("123"), MagicMock(), 5, 5) + + assert sleep_calls == [5, 10] + + +def test_reconnect_loop_resets_backoff_after_stable_connection(): + outcomes = iter( + [ + ("error", None), + ("error", None), + ("error", None), + ("disconnected", 401), + ], + ) + sleep_calls = [] + # start1=0 end1=0 (unstable) | start2=10 end2=10 (unstable) | start3=20 end3=60 (stable, resets) | start4=unused + monotonic_values = iter([0, 0, 10, 10, 20, 60, 999]) + + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=lambda *a: next(outcomes), + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.sleep", + side_effect=sleep_calls.append, + ), patch( + "aikido_zen.background_process.realtime.sse_client.time.monotonic", + side_effect=lambda: next(monotonic_values), + ), patch( + "aikido_zen.background_process.realtime.sse_client.random.random", + return_value=0, + ): + _reconnect_loop(Token("123"), MagicMock(), 5, 5) + + assert sleep_calls == [5, 10, 5] + + +def test_reconnect_loop_logs_unexpected_exception(caplog): + with patch( + "aikido_zen.background_process.realtime.sse_client._connect", + side_effect=RuntimeError("boom"), + ): + _reconnect_loop(Token("123"), MagicMock(), 5, 5) + + assert "SSE loop error : boom" in caplog.text + + +def test_connect_to_sse_starts_a_daemon_thread_running_the_reconnect_loop(): + with patch( + "aikido_zen.background_process.realtime.sse_client._reconnect_loop" + ) as mock_loop: + token = Token("123") + on_event = MagicMock() + thread = connect_to_sse(token, on_event) + thread.join(timeout=2) + + assert thread.daemon is True + mock_loop.assert_called_once_with( + token, on_event, INITIAL_RECONNECT_SECS, READ_TIMEOUT_SECS + ) diff --git a/aikido_zen/background_process/realtime/sse_client/parser.py b/aikido_zen/background_process/realtime/sse_client/parser.py new file mode 100644 index 000000000..783370956 --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/parser.py @@ -0,0 +1,127 @@ +""" +Based on https://github.com/mpetazzoni/sseclient +Licensed under the Apache License, Version 2.0 +Copyright Maxime Petazzoni + +Server Side Events (SSE) client for Python. + +Provides a generator of SSE received through an existing HTTP response. +""" + +_FIELD_SEPARATOR = ":" + + +class SSEParser: + """Implementation of a SSE parser. + + See http://www.w3.org/TR/2009/WD-eventsource-20091029/ for the + specification. + """ + + def __init__(self, event_source, char_enc="utf-8"): + """Initialize the SSE parser over an existing, ready to consume + event source. + + The event source is expected to be a binary stream and have a close() + method. That would usually be something that implements + io.BinaryIOBase, like an httplib or urllib3 HTTPResponse object. + """ + self._event_source = event_source + self._char_enc = char_enc + + def _read(self): + """Read the incoming event source stream and yield event chunks. + + Unfortunately it is possible for some servers to decide to break an + event into multiple HTTP chunks in the response. It is thus necessary + to correctly stitch together consecutive response chunks and find the + SSE delimiter (empty new line) to yield full, correct event chunks.""" + data = b"" + for chunk in self._event_source: + for line in chunk.splitlines(True): + data += line + if data.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")): + yield data + data = b"" + if data: + yield data + + def events(self): + for chunk in self._read(): + event = Event() + # Split before decoding so splitlines() only uses \r and \n + for line in chunk.splitlines(): + # Decode the line. + line = line.decode(self._char_enc) + + # Lines starting with a separator are comments and are to be + # ignored. + if not line.strip() or line.startswith(_FIELD_SEPARATOR): + continue + + data = line.split(_FIELD_SEPARATOR, 1) + field = data[0] + + # Ignore unknown fields. + if field not in event.__dict__: + continue + + if len(data) > 1: + # From the spec: + # "If value starts with a single U+0020 SPACE character, + # remove it from value." + if data[1].startswith(" "): + value = data[1][1:] + else: + value = data[1] + else: + # If no value is present after the separator, + # assume an empty value. + value = "" + + # The data field may come over multiple lines and their values + # are concatenated with each other. + if field == "data": + event.__dict__[field] += value + "\n" + else: + event.__dict__[field] = value + + # Events with no data are not dispatched. + if not event.data: + continue + + # If the data field ends with a newline, remove it. + if event.data.endswith("\n"): + event.data = event.data[0:-1] + + # Empty event names default to 'message' + event.event = event.event or "message" + + # Dispatch the event + yield event + + def close(self): + """Manually close the event source stream.""" + self._event_source.close() + + +class Event: + """Representation of an event from the event stream.""" + + def __init__(self, event_id=None, event="message", data="", retry=None): + self.id = event_id + self.event = event + self.data = data + self.retry = retry + + def __str__(self): + s = f"{self.event} event" + if self.id: + s += f" #{self.id}" + if self.data: + s += f", {len(self.data)} byte{'s' if len(self.data) else ''}" + else: + s += ", no data" + if self.retry: + s += f", retry in {self.retry}ms" + return s diff --git a/aikido_zen/background_process/realtime/sse_client/parser_test.py b/aikido_zen/background_process/realtime/sse_client/parser_test.py new file mode 100644 index 000000000..e9765c93e --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/parser_test.py @@ -0,0 +1,186 @@ +from unittest.mock import MagicMock +from .parser import SSEParser, Event + + +def make_source(chunks): + """Creates an iterable of byte chunks with a close() method, like an HTTP response.""" + source = MagicMock() + source.__iter__.return_value = iter(chunks) + return source + + +def test_read_single_chunk_single_event(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\n\n"] + + +def test_read_splits_multiple_events_in_one_chunk(): + source = make_source([b"data: one\n\ndata: two\n\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: one\n\n", b"data: two\n\n"] + + +def test_read_stitches_event_split_across_chunks(): + source = make_source([b"data: hel", b"lo\n\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\n\n"] + + +def test_read_handles_carriage_return_delimiters(): + source = make_source([b"data: hello\r\r"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\r\r"] + + +def test_read_handles_crlf_delimiters(): + source = make_source([b"data: hello\r\n\r\n"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello\r\n\r\n"] + + +def test_read_yields_trailing_incomplete_chunk(): + source = make_source([b"data: hello"]) + parser = SSEParser(source) + + assert list(parser._read()) == [b"data: hello"] + + +def test_events_basic_data_field(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + events = list(parser.events()) + + assert len(events) == 1 + assert events[0].data == "hello" + assert events[0].event == "message" + + +def test_events_defaults_event_name_to_message(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.event == "message" + + +def test_events_custom_event_name_and_id(): + source = make_source([b"event: update\nid: 42\ndata: payload\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.event == "update" + assert event.id == "42" + assert event.data == "payload" + + +def test_events_multiline_data_is_joined_with_newlines(): + source = make_source([b"data: line1\ndata: line2\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "line1\nline2" + + +def test_events_ignores_comment_lines(): + source = make_source([b": this is a comment\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "hello" + + +def test_events_ignores_unknown_fields(): + source = make_source([b"foo: bar\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "hello" + assert not hasattr(event, "foo") + + +def test_events_strips_single_leading_space_from_value(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == "hello" + + +def test_events_keeps_extra_leading_spaces_beyond_first(): + source = make_source([b"data: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.data == " hello" + + +def test_events_field_with_no_value_is_empty_string(): + source = make_source([b"event\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.event == "message" + + +def test_events_skips_events_with_no_data(): + source = make_source([b"event: ping\n\ndata: hello\n\n"]) + parser = SSEParser(source) + + events = list(parser.events()) + assert len(events) == 1 + assert events[0].data == "hello" + + +def test_events_multiple_events_in_stream(): + source = make_source([b"data: first\n\ndata: second\n\n"]) + parser = SSEParser(source) + + events = list(parser.events()) + assert [e.data for e in events] == ["first", "second"] + + +def test_events_retry_field(): + source = make_source([b"retry: 5000\ndata: hello\n\n"]) + parser = SSEParser(source) + + [event] = parser.events() + assert event.retry == "5000" + + +def test_events_uses_configured_char_encoding(): + source = make_source(["data: héllo\n\n".encode("latin-1")]) + parser = SSEParser(source, char_enc="latin-1") + + [event] = parser.events() + assert event.data == "héllo" + + +def test_close_delegates_to_event_source(): + source = make_source([]) + parser = SSEParser(source) + + parser.close() + + source.close.assert_called_once() + + +def test_event_str_with_data_and_id(): + event = Event(event_id="1", event="message", data="hello") + assert str(event) == "message event #1, 5 bytes" + + +def test_event_str_without_data(): + event = Event() + assert str(event) == "message event, no data" + + +def test_event_str_with_retry(): + event = Event(data="hi", retry=3000) + assert str(event) == "message event, 2 bytes, retry in 3000ms" From cfabb9d98bd62d5980985973745b91a7c0c9d417 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Thu, 16 Jul 2026 13:25:04 +0200 Subject: [PATCH 2/4] Extend test coverage --- .../realtime/sse_client/real_server_test.py | 216 ++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 aikido_zen/background_process/realtime/sse_client/real_server_test.py diff --git a/aikido_zen/background_process/realtime/sse_client/real_server_test.py b/aikido_zen/background_process/realtime/sse_client/real_server_test.py new file mode 100644 index 000000000..4fae2e880 --- /dev/null +++ b/aikido_zen/background_process/realtime/sse_client/real_server_test.py @@ -0,0 +1,216 @@ +""" +Integration tests that exercise the SSE client against a real HTTP server +(as opposed to init_test.py, which mocks urllib.request.urlopen), so that real +socket/timeout/connection-refused behavior is covered, not just our own mocks. +""" + +import http.server +import json +import threading +import time +from unittest.mock import patch + +from aikido_zen.helpers.token import Token +from . import _connect, _reconnect_loop + +REALTIME_URL_PATCH_TARGET = "aikido_zen.background_process.realtime.get_realtime_url" + + +class _TestServer(http.server.ThreadingHTTPServer): + daemon_threads = True # Don't let lingering handler threads block process exit + + +def make_handler(do_get): + """Builds a BaseHTTPRequestHandler class that delegates GET handling to `do_get(handler)`.""" + + class Handler(http.server.BaseHTTPRequestHandler): + protocol_version = "HTTP/1.0" + + def do_GET(self): # pylint: disable=invalid-name + do_get(self) + + def log_message(self, *args, **kwargs): + pass + + return Handler + + +def start_server(handler_class): + server = _TestServer(("127.0.0.1", 0), handler_class) + threading.Thread(target=server.serve_forever, daemon=True).start() + return server, f"http://127.0.0.1:{server.server_address[1]}/" + + +def stop_server(server): + server.shutdown() + server.server_close() + + +def wait_until(predicate, timeout=2.0, interval=0.01): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(interval) + return predicate() + + +def test_happy_path_receives_events_with_auth_header(): + captured = {} + events = [] + + def do_get(handler): + captured["auth"] = handler.headers.get("Authorization") + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + handler.wfile.write(b": ping\n\n") + handler.wfile.write( + b'event: config-updated\ndata: {"configUpdatedAt": 100}\n\n' + ) + handler.wfile.flush() + time.sleep(0.2) + handler.wfile.write( + b'event: config-updated\ndata: {"configUpdatedAt": 200}\n\n' + ) + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + outcome, status_code = _connect( + Token("my-secret-token"), events.append, read_timeout_secs=5 + ) + finally: + stop_server(server) + + assert outcome == "disconnected" + assert status_code == 200 + assert captured["auth"] == "my-secret-token" + assert len(events) == 2 + assert events[0].event == "config-updated" + assert json.loads(events[0].data) == {"configUpdatedAt": 100} + assert json.loads(events[1].data) == {"configUpdatedAt": 200} + + +def test_stops_reconnecting_on_401(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + handler.send_response(401) + handler.end_headers() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + _reconnect_loop( + Token("bad-token"), + lambda event: None, + initial_reconnect_secs=0.05, + read_timeout_secs=1, + ) + finally: + stop_server(server) + + assert connection_count["n"] == 1 + + +def test_reconnects_after_non_200_status(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + if connection_count["n"] == 1: + handler.send_response(500) + handler.end_headers() + return + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + handler.wfile.write(b": ping\n\n") + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + thread = threading.Thread( + target=_reconnect_loop, + args=(Token("test-token"), lambda event: None, 0.05, 1), + daemon=True, + ) + thread.start() + assert wait_until(lambda: connection_count["n"] >= 2) + finally: + stop_server(server) + + +def test_connection_refused(): + probe = _TestServer(("127.0.0.1", 0), make_handler(lambda handler: None)) + port = probe.server_address[1] + probe.server_close() + + with patch(REALTIME_URL_PATCH_TARGET, return_value=f"http://127.0.0.1:{port}/"): + outcome, status_code = _connect( + Token("test-token"), lambda event: None, read_timeout_secs=1 + ) + + assert outcome == "error" + assert status_code is None + + +def test_reconnects_on_read_timeout(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + if connection_count["n"] == 1: + time.sleep(1) + else: + handler.wfile.write(b": ping\n\n") + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + thread = threading.Thread( + target=_reconnect_loop, + args=(Token("test-token"), lambda event: None, 0.05, 0.2), + daemon=True, + ) + thread.start() + assert wait_until(lambda: connection_count["n"] >= 2) + finally: + stop_server(server) + + +def test_reconnects_when_server_closes_connection(): + connection_count = {"n": 0} + + def do_get(handler): + connection_count["n"] += 1 + handler.send_response(200) + handler.send_header("Content-Type", "text/event-stream") + handler.send_header("Cache-Control", "no-cache") + handler.end_headers() + handler.wfile.write(b": ping\n\n") + handler.wfile.flush() + + server, url = start_server(make_handler(do_get)) + try: + with patch(REALTIME_URL_PATCH_TARGET, return_value=url): + thread = threading.Thread( + target=_reconnect_loop, + args=(Token("test-token"), lambda event: None, 0.05, 1), + daemon=True, + ) + thread.start() + assert wait_until(lambda: connection_count["n"] >= 2) + finally: + stop_server(server) From 024c641220244e85caf13e3b2bd86cca8bb10b95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Thu, 16 Jul 2026 13:47:16 +0200 Subject: [PATCH 3/4] Fix Codequality comments --- .../realtime/sse_client/parser.py | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/aikido_zen/background_process/realtime/sse_client/parser.py b/aikido_zen/background_process/realtime/sse_client/parser.py index 783370956..c6e48c607 100644 --- a/aikido_zen/background_process/realtime/sse_client/parser.py +++ b/aikido_zen/background_process/realtime/sse_client/parser.py @@ -36,15 +36,18 @@ def _read(self): event into multiple HTTP chunks in the response. It is thus necessary to correctly stitch together consecutive response chunks and find the SSE delimiter (empty new line) to yield full, correct event chunks.""" - data = b"" + buffer = [] + tail = b"" for chunk in self._event_source: for line in chunk.splitlines(True): - data += line - if data.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")): - yield data - data = b"" - if data: - yield data + buffer.append(line) + tail = (tail + line)[-4:] + if tail.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")): + yield b"".join(buffer) + buffer = [] + tail = b"" + if buffer: + yield b"".join(buffer) def events(self): for chunk in self._read(): @@ -54,8 +57,7 @@ def events(self): # Decode the line. line = line.decode(self._char_enc) - # Lines starting with a separator are comments and are to be - # ignored. + # SSE spec: skip empty lines and colon-prefixed comment lines. if not line.strip() or line.startswith(_FIELD_SEPARATOR): continue @@ -119,7 +121,7 @@ def __str__(self): if self.id: s += f" #{self.id}" if self.data: - s += f", {len(self.data)} byte{'s' if len(self.data) else ''}" + s += f", {len(self.data)} byte{'s' if len(self.data) != 1 else ''}" else: s += ", no data" if self.retry: From 63de88ab3ae6a308ef552e1d9938413d6b5a76bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Fri, 17 Jul 2026 09:23:23 +0200 Subject: [PATCH 4/4] Send X-Agent headers to realtime --- .../background_process/realtime/sse_client/__init__.py | 3 +++ .../realtime/sse_client/real_server_test.py | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/aikido_zen/background_process/realtime/sse_client/__init__.py b/aikido_zen/background_process/realtime/sse_client/__init__.py index aa0ee10aa..d46d3bd81 100644 --- a/aikido_zen/background_process/realtime/sse_client/__init__.py +++ b/aikido_zen/background_process/realtime/sse_client/__init__.py @@ -15,6 +15,7 @@ import aikido_zen.background_process.realtime as realtime from aikido_zen.helpers.logging import logger from .parser import SSEParser +import aikido_zen.config as config INITIAL_RECONNECT_SECS = 5 MAX_RECONNECT_SECS = 60 @@ -94,6 +95,8 @@ def _connect(token, on_event, read_timeout_secs): "Authorization": str(token), "Accept": "text/event-stream", "Cache-Control": "no-cache", + "X-Agent-Platform": "python", + "X-Agent-Version": config.PKG_VERSION, }, ) diff --git a/aikido_zen/background_process/realtime/sse_client/real_server_test.py b/aikido_zen/background_process/realtime/sse_client/real_server_test.py index 4fae2e880..7a216879a 100644 --- a/aikido_zen/background_process/realtime/sse_client/real_server_test.py +++ b/aikido_zen/background_process/realtime/sse_client/real_server_test.py @@ -61,6 +61,10 @@ def test_happy_path_receives_events_with_auth_header(): def do_get(handler): captured["auth"] = handler.headers.get("Authorization") + captured["accept"] = handler.headers.get("Accept") + captured["cache"] = handler.headers.get("Cache-Control") + captured["x-agent-platform"] = handler.headers.get("X-Agent-Platform") + captured["x-agent-version"] = handler.headers.get("X-Agent-Version") handler.send_response(200) handler.send_header("Content-Type", "text/event-stream") handler.send_header("Cache-Control", "no-cache") @@ -88,6 +92,10 @@ def do_get(handler): assert outcome == "disconnected" assert status_code == 200 assert captured["auth"] == "my-secret-token" + assert captured["accept"] == "text/event-stream" + assert captured["cache"] == "no-cache" + assert captured["x-agent-platform"] == "python" + assert captured["x-agent-version"] == "1.0-REPLACE-VERSION" assert len(events) == 2 assert events[0].event == "config-updated" assert json.loads(events[0].data) == {"configUpdatedAt": 100}