fix(python-sdk): wait for in-flight cleanup in dispose() (#2892)

Author: harbournick

packages/sdk/langs/python/superdoc/transport.py

@@ -435,24 +435,33 @@ async def connect(self) -> None:
         """Ensure the host process is running and handshake is complete."""
         await self._ensure_connected()
 
+    async def _await_in_flight_cleanup(self) -> None:
+        """Wait for an already-running cleanup task, if any."""
+        existing = self._cleanup_task
+        if existing and not existing.done():
+            try:
+                await asyncio.shield(existing)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                pass
+
     async def dispose(self) -> None:
         """Gracefully shut down the host process."""
         if self._state == _State.DISCONNECTED:
+            # Reader-triggered cleanup flips state to DISCONNECTED before the
+            # subprocess is fully reaped. If that cleanup is still in flight,
+            # wait for it so dispose() doesn't return while the host process
+            # is still being torn down.
+            await self._await_in_flight_cleanup()
             return
         if self._state == _State.DISPOSING:
             # A reader-triggered cleanup is in flight (or an earlier teardown
             # left state in DISPOSING briefly). Wait for it so the caller
             # observes "host fully torn down" by the time dispose() returns.
             # shield() so a cancelled dispose() doesn't interrupt _cleanup
             # mid-flight and leak the host process.
-            existing = self._cleanup_task
-            if existing and not existing.done():
-                try:
-                    await asyncio.shield(existing)
-                except asyncio.CancelledError:
-                    raise
-                except Exception:
-                    pass
+            await self._await_in_flight_cleanup()
             return
 
         self._stopping = True

packages/sdk/langs/python/tests/test_transport.py

@@ -764,6 +764,54 @@ async def slow_cleanup(error):
         finally:
             _cleanup_wrapper(cli)
 
+    @pytest.mark.asyncio
+    async def test_dispose_waits_for_cleanup_after_state_flips_disconnected(self):
+        # `_cleanup()` flips state to DISCONNECTED before awaiting
+        # `process.wait()`. dispose() must still wait for that cleanup task
+        # instead of short-circuiting and returning while teardown is in
+        # flight.
+        cli = _mock_cli_bin({'handshake': 'ok'})
+        try:
+            transport = AsyncHostTransport(cli, startup_timeout_ms=5_000)
+            await transport.connect()
+            process = transport._process
+            assert process is not None
+
+            wait_started = asyncio.Event()
+            release = asyncio.Event()
+            real_wait = process.wait
+
+            async def slow_wait():
+                wait_started.set()
+                await release.wait()
+                return await real_wait()
+
+            process.wait = slow_wait  # type: ignore[assignment]
+
+            transport._schedule_cleanup(
+                SuperDocError('reader-overflow', code=HOST_PROTOCOL_ERROR),
+            )
+            cleanup_task = transport._cleanup_task
+            assert cleanup_task is not None
+
+            await asyncio.wait_for(wait_started.wait(), timeout=2.0)
+            assert transport.state == 'DISCONNECTED'
+            assert transport._process is None
+            assert not cleanup_task.done()
+
+            dispose_task = asyncio.create_task(transport.dispose())
+            await asyncio.sleep(0.05)
+            assert not dispose_task.done()
+
+            release.set()
+            await dispose_task
+            assert transport.state == 'DISCONNECTED'
+            assert transport._cleanup_task is None
+            await process.wait()
+            assert process.returncode is not None
+        finally:
+            _cleanup_wrapper(cli)
+
     @pytest.mark.asyncio
     async def test_ensure_connected_drains_in_flight_cleanup_before_spawn(self):
         # Round-3 regression: without this drain, `_start_host` reassigns