raise the datastream WebSocket message size limit (SCH-7098)

src/schematic/client.py

@@ -109,6 +109,8 @@ class DataStreamConfig:
     replicator_mode: bool = False
     replicator_health_url: Optional[str] = None
     replicator_health_check: Optional[int] = None
+    # Largest WebSocket message in bytes we accept; leave unset for the default.
+    max_message_size: Optional[int] = None
 
 
 @dataclass
@@ -507,6 +509,8 @@ def __init__(self, api_key: str, config: Optional[AsyncSchematicConfig] = None):
                 ds_opts.replicator_health_url = ds.replicator_health_url
             if ds.replicator_health_check is not None:
                 ds_opts.replicator_health_check = ds.replicator_health_check
+            if ds.max_message_size is not None:
+                ds_opts.max_message_size = ds.max_message_size
 
             self._datastream_client = DataStreamClient(ds_opts)
 

src/schematic/datastream/datastream_client.py

@@ -17,7 +17,7 @@
 from .merge import partial_company, partial_user
 from .rules_engine import RulesEngineClient
 from .types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType, KeyConflictError, MessageType, RulesEngineError
-from .websocket_client import ClientOptions as WSClientOptions, DatastreamWSClient
+from .websocket_client import MAX_MESSAGE_SIZE, ClientOptions as WSClientOptions, DatastreamWSClient
 
 
 _hints_cache: Dict[type, Dict[str, Any]] = {}
@@ -103,6 +103,9 @@ class DataStreamClientOptions:
     replicator_health_url: Optional[str] = "http://localhost:8090/ready"
     replicator_health_check: int = DEFAULT_REPLICATOR_HEALTH_CHECK_MS
 
+    # Largest WebSocket message in bytes we accept; None removes the limit.
+    max_message_size: Optional[int] = MAX_MESSAGE_SIZE
+
     # Event callbacks
     on_connected: Optional[Callable[[], None]] = None
     on_disconnected: Optional[Callable[[], None]] = None
@@ -139,6 +142,7 @@ def __init__(self, options: DataStreamClientOptions) -> None:
         self._base_url = options.base_url
         self._logger = options.logger
         self._cache_ttl = options.cache_ttl
+        self._max_message_size = options.max_message_size
 
         # Callbacks
         self._on_connected = options.on_connected
@@ -237,6 +241,7 @@ async def start(self) -> None:
             message_handler=self._handle_message,
             logger=self._logger,
             connection_ready_handler=self._handle_connection_ready,
+            max_message_size=self._max_message_size,
             on_connected=self._on_ws_connected,
             on_disconnected=self._on_ws_disconnected,
             on_ready=self._on_ready,

src/schematic/datastream/websocket_client.py

@@ -40,6 +40,18 @@
 MAX_RECONNECT_ATTEMPTS = 10
 MIN_RECONNECT_DELAY = 1.0  # seconds
 MAX_RECONNECT_DELAY = 30.0  # seconds
+# A connection only counts as healthy — and so only clears the backoff — once it
+# has delivered a message or stayed up this long. Clearing on the handshake
+# alone means a failure that always lands *after* connecting (an oversized
+# frame, say) reconnects at the backoff floor forever instead of escalating.
+HEALTHY_CONNECTION_THRESHOLD = 30.0  # seconds
+
+# Largest message we accept from the server. The websockets library defaults to
+# 1 MiB, which the flags payload of a large environment exceeds; the library
+# then closes the connection with a 1009 and we reconnect into the same frame.
+# 100 MiB matches the Node SDK's `ws` default and still backstops a runaway
+# payload.
+MAX_MESSAGE_SIZE = 100 * 1024 * 1024  # 100 MiB
 
 # Headers attached to the WebSocket handshake so the backend can distinguish
 # direct-SDK connections from the schematic-datastream-replicator and correlate
@@ -107,6 +119,8 @@ class ClientOptions:
     max_reconnect_attempts: int = MAX_RECONNECT_ATTEMPTS
     min_reconnect_delay: float = MIN_RECONNECT_DELAY
     max_reconnect_delay: float = MAX_RECONNECT_DELAY
+    # Maximum size in bytes of a message we accept; None removes the limit.
+    max_message_size: Optional[int] = MAX_MESSAGE_SIZE
 
     # Event callbacks — called on state transitions
     on_connected: Optional[Callable[[], None]] = None
@@ -165,6 +179,7 @@ def __init__(self, options: ClientOptions) -> None:
         self._max_reconnect_attempts = options.max_reconnect_attempts
         self._min_reconnect_delay = options.min_reconnect_delay
         self._max_reconnect_delay = options.max_reconnect_delay
+        self._max_message_size = options.max_message_size
 
         # Event callbacks
         self._on_connected = options.on_connected
@@ -268,9 +283,12 @@ async def _connect_and_read(self) -> None:
                     # avoid conflicts with the Go server's gorilla/websocket.
                     ping_interval=None,
                     ping_timeout=None,
+                    # Without this the library caps messages at 1 MiB and closes
+                    # the connection with a 1009 on anything larger.
+                    max_size=self._max_message_size,
                 ) as ws:
                     self._ws = ws
-                    self._reconnect_attempts = 0
+                    connected_at = asyncio.get_event_loop().time()
                     self._set_connected(True)
 
                     # Run the ready handler before marking the client ready
@@ -295,8 +313,16 @@ async def _connect_and_read(self) -> None:
                     try:
                         async for raw_message in ws:
                             await self._handle_message(raw_message)
+                            # A delivered message proves the connection works,
+                            # so start the next backoff from scratch.
+                            self._reconnect_attempts = 0
                     finally:
                         self._stop_ping_pong()
+                        # A long-lived connection counts as healthy too, even if
+                        # the server never sent anything.
+                        elapsed = asyncio.get_event_loop().time() - connected_at
+                        if elapsed >= HEALTHY_CONNECTION_THRESHOLD:
+                            self._reconnect_attempts = 0
 
                     self._logger.info("WebSocket connection closed")
 

tests/datastream/test_websocket_client.py

@@ -12,6 +12,7 @@
 from schematic.datastream.types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType
 from schematic.datastream.websocket_client import (
     _WS_HEADERS_KWARG,
+    MAX_MESSAGE_SIZE,
     ClientOptions,
     DatastreamWSClient,
     convert_api_url_to_websocket_url,
@@ -627,3 +628,213 @@ async def handler(msg): pass
     # At high attempt counts, delay should be capped at max + jitter ceiling
     delay = client._calculate_backoff_delay(20)
     assert delay <= 5.0 + 1.0
+
+
+# ---------------------------------------------------------------------------
+# Message size limit
+# ---------------------------------------------------------------------------
+
+
+async def test_max_size_passed_to_websockets_connect() -> None:
+    """The connect call carries an explicit max_size, so we don't inherit the
+    websockets library's 1 MiB default."""
+    captured_kwargs: dict = {}
+    ws = MockWebSocket(block_on_empty=True)
+
+    @asynccontextmanager
+    async def capturing_connect(*args, **kwargs):
+        captured_kwargs.update(kwargs)
+        yield ws
+
+    connected = asyncio.Event()
+    client, ws, _ = make_client(ws=ws, on_connected=lambda: connected.set())
+
+    with patch("schematic.datastream.websocket_client.websockets.connect", capturing_connect):
+        async with run_client(client):
+            await asyncio.wait_for(connected.wait(), timeout=2.0)
+
+    assert captured_kwargs["max_size"] == MAX_MESSAGE_SIZE
+    assert MAX_MESSAGE_SIZE > 1024 * 1024
+
+
+async def test_max_message_size_is_configurable() -> None:
+    captured_kwargs: dict = {}
+    ws = MockWebSocket(block_on_empty=True)
+
+    @asynccontextmanager
+    async def capturing_connect(*args, **kwargs):
+        captured_kwargs.update(kwargs)
+        yield ws
+
+    connected = asyncio.Event()
+    client, ws, _ = make_client(ws=ws, max_message_size=None, on_connected=lambda: connected.set())
+
+    with patch("schematic.datastream.websocket_client.websockets.connect", capturing_connect):
+        async with run_client(client):
+            await asyncio.wait_for(connected.wait(), timeout=2.0)
+
+    assert captured_kwargs["max_size"] is None
+
+
+@asynccontextmanager
+async def serve_one_message(payload: str) -> AsyncIterator[str]:
+    """Run a real websockets server that sends `payload` to each client.
+
+    Yields the ws:// URL to connect to.
+    """
+    import websockets as ws_lib
+
+    async def handler(connection, *_args) -> None:
+        await connection.send(payload)
+        try:
+            await connection.wait_closed()
+        except Exception:
+            pass
+
+    server = await ws_lib.serve(handler, "127.0.0.1", 0)
+    try:
+        port = next(iter(server.sockets)).getsockname()[1]
+        yield f"ws://127.0.0.1:{port}/datastream"
+    finally:
+        server.close()
+        await server.wait_closed()
+
+
+def _oversized_payload() -> str:
+    """A datastream message larger than the websockets library's 1 MiB default."""
+    message = json.dumps(
+        {
+            "entity_type": "rulesengine.Flags",
+            "message_type": "full",
+            "data": {"filler": "x" * (2 * 1024 * 1024)},
+        }
+    )
+    assert len(message.encode()) > 1024 * 1024
+    return message
+
+
+async def test_oversized_message_reaches_handler_over_a_real_connection() -> None:
+    """Regression for SCH-7098: a >1 MiB Flags frame used to be dropped by the
+    websockets library with a 1009, leaving the client in a reconnect loop."""
+    payload = _oversized_payload()
+    received: List[DataStreamResp] = []
+
+    async def handler(m: DataStreamResp) -> None:
+        received.append(m)
+
+    async with serve_one_message(payload) as url:
+        client = DatastreamWSClient(
+            ClientOptions(
+                url=url,
+                api_key="key",
+                message_handler=handler,
+                logger=logger,
+                min_reconnect_delay=0.0,
+                max_reconnect_delay=0.0,
+            )
+        )
+        async with run_client(client):
+            await wait_until(lambda: len(received) == 1, timeout=10.0)
+
+    assert received[0].entity_type == "rulesengine.Flags"
+    assert len(received[0].data["filler"]) == 2 * 1024 * 1024  # type: ignore[index]
+
+
+async def test_oversized_message_is_dropped_at_the_old_1mib_limit() -> None:
+    """Control for the test above: pin max_message_size back to the library
+    default and the same frame never reaches the handler."""
+    payload = _oversized_payload()
+    received: List[DataStreamResp] = []
+    connects: List[int] = []
+
+    async def handler(m: DataStreamResp) -> None:
+        received.append(m)
+
+    async with serve_one_message(payload) as url:
+        client = DatastreamWSClient(
+            ClientOptions(
+                url=url,
+                api_key="key",
+                message_handler=handler,
+                logger=logger,
+                min_reconnect_delay=0.0,
+                max_reconnect_delay=0.0,
+                max_message_size=1024 * 1024,
+                on_connected=lambda: connects.append(1),
+            )
+        )
+        async with run_client(client):
+            await wait_until(lambda: len(connects) >= 2, timeout=10.0)
+
+    assert received == []
+
+
+# ---------------------------------------------------------------------------
+# Backoff escalation
+# ---------------------------------------------------------------------------
+
+
+async def test_backoff_escalates_when_failure_follows_a_successful_connect() -> None:
+    """A failure that always lands after the handshake must still escalate the
+    backoff — otherwise a poison frame reconnects at the floor forever."""
+    attempts_seen: List[int] = []
+
+    @asynccontextmanager
+    async def connect_then_fail(*args, **kwargs):
+        attempts_seen.append(0)
+        yield MockWebSocket()
+        raise ConnectionError("closed right after connecting")
+
+    async def handler(m: DataStreamResp) -> None: ...
+
+    client = DatastreamWSClient(
+        ClientOptions(
+            url="wss://test.example.com/datastream",
+            api_key="key",
+            message_handler=handler,
+            logger=logger,
+            min_reconnect_delay=0.0,
+            max_reconnect_delay=0.0,
+            max_reconnect_attempts=4,
+        )
+    )
+
+    with patch("schematic.datastream.websocket_client.websockets.connect", connect_then_fail):
+        async with run_client(client):
+            await wait_until(lambda: len(attempts_seen) >= 3, timeout=5.0)
+            assert client._reconnect_attempts >= 2
+
+
+async def test_delivered_message_clears_the_backoff() -> None:
+    """A connection that delivers a message is healthy, so the next failure
+    starts its backoff from zero again."""
+    msg = json.dumps({"entity_type": "rulesengine.Company", "message_type": "full", "data": {"id": "c1"}})
+    received: List[DataStreamResp] = []
+    connects: List[int] = []
+
+    async def handler(m: DataStreamResp) -> None:
+        received.append(m)
+
+    @asynccontextmanager
+    async def connect(*args, **kwargs):
+        connects.append(1)
+        if len(connects) == 1:
+            raise ConnectionError("first attempt fails")
+        yield MockWebSocket(messages=[msg], block_on_empty=True)
+
+    client = DatastreamWSClient(
+        ClientOptions(
+            url="wss://test.example.com/datastream",
+            api_key="key",
+            message_handler=handler,
+            logger=logger,
+            min_reconnect_delay=0.0,
+            max_reconnect_delay=0.0,
+            max_reconnect_attempts=5,
+        )
+    )
+
+    with patch("schematic.datastream.websocket_client.websockets.connect", connect):
+        async with run_client(client):
+            await wait_until(lambda: len(received) == 1)
+            assert client._reconnect_attempts == 0