fix: exporting a large figure hangs

Choreographer uses remote debugging pipes to execute Javascript
functions in Chrome. These have a receive buffer size of 100MB.
Previously, if a spec exceeded this size, we would try to send it
anyways and the image export would never complete. This makes it
so we chunk large specs to keep each send within the buffer size.

chrome ref: https://chromium.googlesource.com/chromium/src/+/refs/heads/main/content/browser/devtools/devtools_pipe_handler.cc#44
issue: #419

Author: Elias Talcott

src/py/kaleido/_kaleido_tab/_tab.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING
 
 import logistro
+import orjson
 
 from . import _devtools_utils as _dtools
 from . import _js_logger
@@ -19,10 +20,18 @@
 
 
 _TEXT_FORMATS = ("svg", "json")  # eps
+_CHUNK_SIZE = 10 * 1024 * 1024  # 10 MB
 
 _logger = logistro.getLogger(__name__)
 
 
+def _orjson_default(obj):
+    """Fallback for types orjson can't handle natively (e.g. NumPy string arrays)."""
+    if hasattr(obj, "tolist"):
+        return obj.tolist()
+    raise TypeError(f"Type is not JSON serializable: {type(obj).__name__}")
+
+
 def _subscribe_new(tab: choreo.Tab, event: str) -> asyncio.Future:
     """Create subscription to tab clearing old ones first: helper function."""
     new_future = tab.subscribe_once(event)
@@ -117,22 +126,38 @@ async def _calc_fig(
         render_prof,
         stepper,
     ) -> bytes:
-        # js script
-        kaleido_js_fn = (
-            r"function(spec, ...args)"
-            r"{"
-            r"return kaleido_scopes.plotly(spec, ...args).then(JSON.stringify);"
-            r"}"
-        )
-        render_prof.profile_log.tick("sending javascript")
-        result = await _dtools.exec_js_fn(
-            self.tab,
-            self._current_js_id,
-            kaleido_js_fn,
+        render_prof.profile_log.tick("serializing spec")
+        spec_str = orjson.dumps(
             spec,
-            topojson,
-            stepper,
-        )
+            default=_orjson_default,
+            option=orjson.OPT_SERIALIZE_NUMPY,
+        ).decode()
+        render_prof.profile_log.tick("spec serialized")
+
+        render_prof.profile_log.tick("sending javascript")
+        if len(spec_str) <= _CHUNK_SIZE:
+            kaleido_js_fn = (
+                r"function(specStr, ...args)"
+                r"{"
+                r"return kaleido_scopes"
+                r".plotly(JSON.parse(specStr), ...args)"
+                r".then(JSON.stringify);"
+                r"}"
+            )
+            result = await _dtools.exec_js_fn(
+                self.tab,
+                self._current_js_id,
+                kaleido_js_fn,
+                spec_str,
+                topojson,
+                stepper,
+            )
+        else:
+            result = await self._calc_fig_chunked(
+                spec_str,
+                topojson=topojson,
+                stepper=stepper,
+            )
         _raise_error(result)
         render_prof.profile_log.tick("javascript sent")
 
@@ -154,3 +179,41 @@ async def _calc_fig(
         render_prof.data_out_size = len(res)
         render_prof.js_log = self.js_logger.log
         return res
+
+    async def _calc_fig_chunked(
+        self,
+        spec_str: str,
+        *,
+        topojson: str | None,
+        stepper,
+    ):
+        await _dtools.exec_js_fn(
+            self.tab,
+            self._current_js_id,
+            r"function() { window.__kaleido_chunks = []; }",
+        )
+
+        for i in range(0, len(spec_str), _CHUNK_SIZE):
+            chunk = spec_str[i : i + _CHUNK_SIZE]
+            await _dtools.exec_js_fn(
+                self.tab,
+                self._current_js_id,
+                r"function(c) { window.__kaleido_chunks.push(c); }",
+                chunk,
+            )
+
+        kaleido_js_fn = (
+            r"function(...args)"
+            r"{"
+            r"var spec = JSON.parse(window.__kaleido_chunks.join(''));"
+            r"delete window.__kaleido_chunks;"
+            r"return kaleido_scopes.plotly(spec, ...args).then(JSON.stringify);"
+            r"}"
+        )
+        return await _dtools.exec_js_fn(
+            self.tab,
+            self._current_js_id,
+            kaleido_js_fn,
+            topojson,
+            stepper,
+        )