diff --git a/packages/sdk/langs/python/superdoc/transport.py b/packages/sdk/langs/python/superdoc/transport.py index d7c4a9ca0e..d7163b7423 100644 --- a/packages/sdk/langs/python/superdoc/transport.py +++ b/packages/sdk/langs/python/superdoc/transport.py @@ -435,9 +435,25 @@ 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 @@ -445,14 +461,7 @@ async def dispose(self) -> None: # 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 diff --git a/packages/sdk/langs/python/tests/test_transport.py b/packages/sdk/langs/python/tests/test_transport.py index f54c0c7a71..c9043968d8 100644 --- a/packages/sdk/langs/python/tests/test_transport.py +++ b/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