fix(async): preserve optimization after result wait timeout

pymilvus/milvus_client/async_optimize_task.py

@@ -84,12 +84,13 @@ def check_cancelled(self) -> None:
             raise MilvusException(message="Optimization task was cancelled")
 
     async def result(self, timeout: Optional[float] = None) -> OptimizeResult:
+        """Wait for the result without cancelling optimization when the wait times out."""
         if not self._task:
             raise MilvusException(message="Task has not been started")
 
         try:
             if timeout is not None:
-                return await asyncio.wait_for(self._task, timeout=timeout)
+                return await asyncio.wait_for(asyncio.shield(self._task), timeout=timeout)
             return await self._task
         except asyncio.TimeoutError as e:
             raise MilvusException(message="Timeout waiting for optimization to complete") from e

tests/unit/test_async_milvus_client_ops.py

@@ -1,3 +1,4 @@
+import asyncio
 import inspect
 from unittest.mock import ANY, AsyncMock, MagicMock, patch
 
@@ -444,6 +445,44 @@ async def test_upsert_exception_propagates(self):
 
 
 class TestAsyncClientOptimize:
+    @pytest.mark.asyncio
+    async def test_wait_timeout_preserves_optimization_and_other_waiter(self):
+        client, handler = _make_client()
+        compacting = asyncio.Event()
+        finish_compaction = asyncio.Event()
+
+        async def compaction_state(*args, **kwargs):
+            compacting.set()
+            await finish_compaction.wait()
+            return MagicMock(state=2)
+
+        handler.compact.return_value = 99
+        handler.get_compaction_state.side_effect = compaction_state
+        handler.get_load_state.return_value = LoadState.Loaded
+        handler._get_schema.return_value = ({"fields": []}, 0)
+        handler.list_indexes.return_value = []
+
+        with patch.object(client, "refresh_load", new_callable=AsyncMock) as refresh_load:
+            task = await client.optimize("col", wait=False)
+            await compacting.wait()
+            other_waiter = asyncio.create_task(task.result())
+            await asyncio.sleep(0)
+            try:
+                with pytest.raises(MilvusException, match="Timeout waiting"):
+                    await task.result(timeout=0)
+                assert not task.done()
+                assert not other_waiter.done()
+                finish_compaction.set()
+                result = await other_waiter
+                assert result.status == "success"
+                assert result.compaction_id == 99
+                assert await task.result() is result
+                handler.compact.assert_awaited_once()
+                refresh_load.assert_awaited_once_with("col", timeout=None)
+            finally:
+                finish_compaction.set()
+                await asyncio.gather(task._task, other_waiter, return_exceptions=True)
+
     @pytest.mark.asyncio
     async def test_is_collection_loaded_true(self):
         client, handler = _make_client()

tests/unit/test_async_optimize_task.py

@@ -63,6 +63,59 @@ def test_check_cancelled_when_cancelled_raises(self):
 
 
 class TestAsyncOptimizeTaskResult:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("timeout", [0, 0.001])
+    @pytest.mark.parametrize("fails", [False, True])
+    async def test_wait_timeout_preserves_background_result(self, timeout, fails):
+        started = asyncio.Event()
+        finish = asyncio.Event()
+        expected = MagicMock(collection_name="col")
+
+        async def execute(**kwargs):
+            started.set()
+            await finish.wait()
+            if fails:
+                raise MilvusException(message="Index rebuild failed")
+            return expected
+
+        task = _make_task(execute)
+        task.start()
+        await started.wait()
+        try:
+            with pytest.raises(MilvusException, match="Timeout waiting"):
+                await task.result(timeout=timeout)
+            assert not task.done()
+            assert not task.cancelled()
+            assert ProgressStage.CANCELLED not in task.progress_history()
+            finish.set()
+            if fails:
+                with pytest.raises(MilvusException, match="Index rebuild failed"):
+                    await task.result()
+            else:
+                assert await task.result() is expected
+        finally:
+            finish.set()
+            await asyncio.gather(task._task, return_exceptions=True)
+
+    @pytest.mark.asyncio
+    async def test_explicit_cancel_interrupts_timed_wait(self):
+        started = asyncio.Event()
+
+        async def execute(**kwargs):
+            started.set()
+            await asyncio.Event().wait()
+
+        task = _make_task(execute)
+        task.start()
+        await started.wait()
+        waiter = asyncio.create_task(task.result(timeout=10))
+        await asyncio.sleep(0)
+        assert task.cancel()
+        with pytest.raises(MilvusException, match="cancelled"):
+            await waiter
+        assert task.done()
+        assert task.cancelled()
+
     def test_result_without_start_raises(self):
         task = _make_task()
 
@@ -84,19 +137,6 @@ async def run():
 
         asyncio.run(run())
 
-    def test_result_timeout_raises(self):
-        async def slow_execute(**kwargs):
-            await asyncio.sleep(10)
-            return MagicMock()
-
-        async def run():
-            task = AsyncOptimizeTask("col", None, None, slow_execute)
-            task.start()
-            with pytest.raises(MilvusException):
-                await task.result(timeout=0.01)
-
-        asyncio.run(run())
-
     def test_result_cancelled_raises(self):
         async def run():
             execute_fn = AsyncMock(return_value=MagicMock())