fix(async): preserve optimization after result wait timeout

pymilvus/milvus_client/async_milvus_client.py

@@ -2,6 +2,7 @@
 import copy
 import time
 import types
+from contextlib import suppress
 from typing import Dict, List, Optional, Type, Union
 
 from pymilvus.client import type_info
@@ -2310,6 +2311,7 @@ async def optimize(
                 If not provided, uses system default.
             wait (bool): Whether to wait for optimization to complete. Defaults to True.
                 If False, returns an OptimizeTask for async tracking.
+                If True, a wait timeout or caller cancellation stops the client-side task.
             timeout (Optional[float]): Maximum time in seconds to wait for optimization.
                 Only applies when wait=True.
             **kwargs: Additional arguments.
@@ -2356,7 +2358,14 @@ async def optimize(
         task.start()
 
         if wait:
-            return await task.result(timeout=timeout)
+            try:
+                return await task.result(timeout=timeout)
+            finally:
+                # No task handle is returned on this path; stop and retrieve unfinished work.
+                if not task.done():
+                    task.cancel()
+                    with suppress(MilvusException):
+                        await task.result()
 
         return task
 

pymilvus/milvus_client/async_optimize_task.py

@@ -84,17 +84,23 @@ def check_cancelled(self) -> None:
             raise MilvusException(message="Optimization task was cancelled")
 
     async def result(self, timeout: Optional[float] = None) -> OptimizeResult:
+        """Wait without cancelling optimization on wait timeout or caller cancellation.
+
+        Call cancel() to stop the background optimization explicitly.
+        """
         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 self._task
+                return await asyncio.wait_for(asyncio.shield(self._task), timeout=timeout)
+            return await asyncio.shield(self._task)
         except asyncio.TimeoutError as e:
             raise MilvusException(message="Timeout waiting for optimization to complete") from e
         except asyncio.CancelledError as e:
-            raise MilvusException(message="Optimization task was cancelled") from e
+            if self._task.cancelled():
+                raise MilvusException(message="Optimization task was cancelled") from e
+            raise
 
     def set_progress(self, stage: ProgressStage) -> None:
         if self._cancelled:

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,97 @@ async def test_upsert_exception_propagates(self):
 
 
 class TestAsyncClientOptimize:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("timeout", [0, 0.001])
+    async def test_wait_true_timeout_cleans_up_owned_task(self, timeout):
+        client, _ = _make_client()
+        tasks = []
+        start = AsyncOptimizeTask.start
+
+        def record_start(task):
+            tasks.append(task)
+            start(task)
+
+        async def execute(**kwargs):
+            await asyncio.Event().wait()
+
+        with patch.object(AsyncOptimizeTask, "start", record_start), patch.object(
+            client, "_execute_optimize", execute
+        ):
+            try:
+                with pytest.raises(MilvusException, match="Timeout waiting"):
+                    await client.optimize("col", wait=True, timeout=timeout)
+                assert len(tasks) == 1
+                assert tasks[0].done()
+                assert tasks[0].cancelled()
+            finally:
+                for task in tasks:
+                    task.cancel()
+                    await asyncio.gather(task._task, return_exceptions=True)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("timeout", [None, 10])
+    async def test_wait_true_caller_cancellation_cleans_up_owned_task(self, timeout):
+        client, _ = _make_client()
+        started = asyncio.Event()
+        tasks = []
+
+        async def execute(task, **kwargs):
+            tasks.append(task)
+            started.set()
+            await asyncio.Event().wait()
+
+        with patch.object(client, "_execute_optimize", execute):
+            waiter = asyncio.create_task(client.optimize("col", wait=True, timeout=timeout))
+            await started.wait()
+            try:
+                waiter.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await waiter
+                assert tasks[0].done()
+                assert tasks[0].cancelled()
+            finally:
+                tasks[0].cancel()
+                await asyncio.gather(waiter, tasks[0]._task, return_exceptions=True)
+
+    @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,143 @@ def test_check_cancelled_when_cancelled_raises(self):
 
 
 class TestAsyncOptimizeTaskResult:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("timeout", [None, 10])
+    async def test_cancelled_result_waiter_preserves_background_task(self, timeout):
+        started = asyncio.Event()
+        finish = asyncio.Event()
+        expected = MagicMock(collection_name="col")
+
+        async def execute(**kwargs):
+            started.set()
+            await finish.wait()
+            return expected
+
+        task = _make_task(execute)
+        task.start()
+        await started.wait()
+        waiter = asyncio.create_task(task.result(timeout=timeout))
+        other_waiter = asyncio.create_task(task.result())
+        await asyncio.sleep(0)
+        try:
+            waiter.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await waiter
+            assert not task.done()
+            assert not task.cancelled()
+            assert not other_waiter.done()
+            finish.set()
+            assert await other_waiter is expected
+        finally:
+            finish.set()
+            await asyncio.gather(task._task, waiter, other_waiter, return_exceptions=True)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("timeout", [None, 10])
+    async def test_external_wait_for_preserves_background_task(self, timeout):
+        started = asyncio.Event()
+        finish = asyncio.Event()
+        expected = MagicMock(collection_name="col")
+
+        async def execute(**kwargs):
+            started.set()
+            await finish.wait()
+            return expected
+
+        task = _make_task(execute)
+        task.start()
+        await started.wait()
+        try:
+            with pytest.raises(asyncio.TimeoutError):
+                await asyncio.wait_for(task.result(timeout=timeout), timeout=0.001)
+            assert not task.done()
+            assert not task.cancelled()
+            finish.set()
+            assert await task.result() is expected
+        finally:
+            finish.set()
+            await asyncio.gather(task._task, return_exceptions=True)
+
+    @pytest.mark.asyncio
+    @pytest.mark.skipif(not hasattr(asyncio, "timeout"), reason="Requires Python 3.11+")
+    async def test_external_timeout_context_preserves_background_task(self):
+        started = asyncio.Event()
+        finish = asyncio.Event()
+        expected = MagicMock(collection_name="col")
+
+        async def execute(**kwargs):
+            started.set()
+            await finish.wait()
+            return expected
+
+        task = _make_task(execute)
+        task.start()
+        await started.wait()
+        try:
+            with pytest.raises(asyncio.TimeoutError):
+                async with asyncio.timeout(0.001):
+                    await task.result()
+            assert not task.done()
+            assert not task.cancelled()
+            finish.set()
+            assert await task.result() is expected
+        finally:
+            finish.set()
+            await asyncio.gather(task._task, return_exceptions=True)
+
+    @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 +221,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())