fix: prevent write_file from failing when streaming data. (#42)

We need to send writeFile data concurrently with rpc call, like in the JS API

Author: marvinhagemeister

src/deno_sandbox/fs.py

@@ -18,7 +18,7 @@
 from re import Pattern
 
 from .process import AbortSignal
-from .stream import stream_data
+from .stream import complete_stream, start_stream
 from .utils import convert_to_camel_case, convert_to_snake_case
 
 if TYPE_CHECKING:
@@ -317,13 +317,13 @@ async def write_file(
             mode: Set the file permission mode.
         """
         if isinstance(data, bytes):
-            # Stream bytes as a single chunk
-            content_stream_id = await stream_data(self._rpc, iter([data]))
+            streamable = iter([data])
         else:
-            # Stream data from iterable/file object
-            content_stream_id = await stream_data(self._rpc, data)
+            streamable = data
 
-        params: dict[str, Any] = {"path": path, "contentStreamId": content_stream_id}
+        stream_id, writer = await start_stream(self._rpc)
+
+        params: dict[str, Any] = {"path": path, "contentStreamId": stream_id}
         options: dict[str, Any] = {}
         if create is not None:
             options["create"] = create
@@ -335,7 +335,20 @@ async def write_file(
             options["mode"] = mode
         if options:
             params["options"] = convert_to_camel_case(options)
-        await self._rpc.call("writeFile", params)
+
+        # Send data concurrently with the RPC call to avoid deadlock:
+        # the server waits for stream data before responding to writeFile.
+        async def _send() -> None:
+            try:
+                await complete_stream(writer, streamable)
+            except Exception as e:
+                await writer.error(str(e))
+
+        task = self._rpc._loop.create_task(_send())
+        try:
+            await self._rpc.call("writeFile", params)
+        finally:
+            await task
 
     async def read_text_file(
         self,

tests/test_stream.py

@@ -1,3 +1,4 @@
+import asyncio
 import io
 import os
 import tempfile
@@ -165,6 +166,26 @@ def large_chunks():
     assert content == b"x" * 102400
 
 
+@pytest.mark.asyncio(loop_scope="session")
+async def test_write_file_slow_async_generator(async_shared_sandbox):
+    """Test write_file with a slow async generator that yields chunks with delays.
+
+    This verifies that stream data is sent concurrently with the RPC call.
+    With incorrect ordering (all data sent before RPC), this would deadlock
+    because the server waits for stream data before responding to writeFile.
+    """
+    sb = async_shared_sandbox
+
+    async def slow_chunks():
+        for i in range(5):
+            await asyncio.sleep(0.05)
+            yield f"chunk{i} ".encode()
+
+    await sb.fs.write_file("test_slow_async.txt", slow_chunks())
+    content = await sb.fs.read_file("test_slow_async.txt")
+    assert content == b"chunk0 chunk1 chunk2 chunk3 chunk4 "
+
+
 # stdin streaming tests for spawn