@camera.ui/server 0.0.105 → 0.0.106

package/dist/server/src/_camera_ui_tools/camera_ui_rpc/chunking.py

@@ -1,9 +1,10 @@
 """Optimized chunking system for large message handling."""
 
+import time
 from collections.abc import Callable, Generator
 from typing import Any
 
-from .codec import decode
+from .codec import decode_message
 from .types import ChunkData
 
 
@@ -152,10 +153,13 @@ class ChunkAssembler:
             raise RuntimeError("Cannot get data: Not all chunks have been received.")
 
         if self.buffer is not None:
-            # Optimal path: The buffer is already fully assembled.
-            # We create a memoryview to avoid copying the buffer before decoding.
+            # Optimal path: The buffer is already fully assembled. Chunked
+            # payloads are always full wire messages, so this goes through
+            # decode_message: a reassembled CUIB frame yields zero-copy
+            # memoryview slices into the assembly buffer (the views keep the
+            # bytearray alive after the assembler is pooled/reset).
             data_view = memoryview(self.buffer)
-            return decode(data_view)
+            return decode_message(data_view)
         else:
             # Fallback path: Manually combine chunks from temporary storage.
             # This path involves a final copy operation.
@@ -167,7 +171,7 @@ class ChunkAssembler:
                 chunk_len = len(chunk)
                 combined[offset : offset + chunk_len] = chunk
                 offset += chunk_len
-            return decode(combined)
+            return decode_message(combined)
 
     def get_progress(self) -> dict[str, int | float]:
         """Get the current progress of the chunk transfer."""
@@ -180,11 +184,22 @@ class ChunkAssembler:
 class ChunkingManager:
     """Manages multiple concurrent chunk transfers."""
 
+    # Incomplete transfers hold their full pre-allocated buffer. If a single
+    # chunk is lost (sender reconnected mid-transfer), the assembler would
+    # otherwise sit in memory until the whole client disconnects — so stale
+    # transfers are swept lazily on every start_receiving/process_chunk call.
+    _STALE_TRANSFER_TTL: float = 30.0
+    # Sweeping on every chunk of a large transfer is pure overhead — one
+    # sweep per second is more than enough for a 30s TTL.
+    _SWEEP_INTERVAL: float = 1.0
+
     def __init__(self) -> None:
         self.assemblers: dict[str, ChunkAssembler] = {}
         self.completed_callbacks: dict[str, Callable[[Any], None]] = {}
         self.error_callbacks: dict[str, Callable[[Exception], None]] = {}
         self._assembler_pool: list[ChunkAssembler] = []
+        self._last_activity: dict[str, float] = {}
+        self._last_sweep: float = 0.0
 
     def _get_assembler(self, transfer_id: str) -> ChunkAssembler:
         """Get a reusable assembler from a pool or create a new one."""
@@ -228,14 +243,42 @@ class ChunkingManager:
             on_error: Callback to execute if an error occurs.
             total_size: The total size of the final data for pre-allocation.
         """
+        self._sweep_stale()
         assembler = self._get_assembler(transfer_id)
         assembler.set_expected_chunks(total_chunks, total_size, chunk_size)
         self.assemblers[transfer_id] = assembler
         self.completed_callbacks[transfer_id] = on_complete
         self.error_callbacks[transfer_id] = on_error
+        self._last_activity[transfer_id] = time.monotonic()
+
+    def _sweep_stale(self) -> None:
+        """Drop incomplete transfers with no activity for _STALE_TRANSFER_TTL.
+
+        Rate-limited to at most one sweep per _SWEEP_INTERVAL — this runs on
+        every start_receiving/process_chunk call.
+        """
+        if not self._last_activity:
+            return
+        now = time.monotonic()
+        if now - self._last_sweep < self._SWEEP_INTERVAL:
+            return
+        self._last_sweep = now
+        cutoff = now - self._STALE_TRANSFER_TTL
+        stale = [tid for tid, ts in self._last_activity.items() if ts < cutoff]
+        for transfer_id in stale:
+            error_callback = self.error_callbacks.pop(transfer_id, None)
+            assembler = self.assemblers.pop(transfer_id, None)
+            if assembler:
+                self._return_assembler(assembler)
+            self.completed_callbacks.pop(transfer_id, None)
+            self._last_activity.pop(transfer_id, None)
+            if error_callback:
+                error_callback(RuntimeError("Transfer timed out"))
 
     def process_chunk(self, chunk: dict[str, Any]) -> None:
         """Process an incoming chunk for a specific transfer."""
+        self._sweep_stale()
+
         transfer_id = chunk.get("id") or chunk.get("transferId")
         if not transfer_id:
             return
@@ -245,6 +288,8 @@ class ChunkingManager:
             # Received a chunk for an unknown or completed transfer.
             return
 
+        self._last_activity[transfer_id] = time.monotonic()
+
         try:
             # Normalize keys in-place to avoid allocating a new dict.
             # Handles minor key differences between JS/Go/Python senders.
@@ -261,6 +306,7 @@ class ChunkingManager:
                 # Cleanup and return assembler to the pool
                 del self.assemblers[transfer_id]
                 self.error_callbacks.pop(transfer_id, None)
+                self._last_activity.pop(transfer_id, None)
                 self._return_assembler(assembler)
 
                 if callback:
@@ -276,6 +322,7 @@ class ChunkingManager:
         if assembler:
             self._return_assembler(assembler)
         self.completed_callbacks.pop(transfer_id, None)
+        self._last_activity.pop(transfer_id, None)
 
         if error_callback:
             error_callback(error)