From eb5c2c989929e8286f588a395effeb3ea5102864 Mon Sep 17 00:00:00 2001
From: yixi ning <ningyixi0@gmail.com>
Date: Sat, 17 Oct 2026 15:22:17 +0000
Subject: [PATCH] fix(core): finish application cleanup before propagating
 cancellation

---
 reme/application.py    |  6 +++++
 tests/unit/test_job.py | 58 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/reme/application.py b/reme/application.py
index 5bbe3d14..4276e721 100644
--- a/reme/application.py
+++ b/reme/application.py
@@ -227,15 +227,21 @@ class Application(BaseComponent):
 
     async def _close_started_components(self) -> None:
         """Close resources while the caller serializes component mutations."""
+        cancelled_error: asyncio.CancelledError | None = None
         for c in reversed(self._started_components):
             try:
                 await c.close()
+            except asyncio.CancelledError as exc:
+                if cancelled_error is None:
+                    cancelled_error = exc
             except Exception as e:
                 self.logger.exception(f"Failed to close {component_type_name(c.component_type)}:{c.name}: {e}")
         self._started_components.clear()
         if self.context.thread_pool is not None:
             self.context.thread_pool.shutdown(wait=True)
             self.context.thread_pool = None
+        if cancelled_error is not None:
+            raise cancelled_error
 
     async def update_component(self, component_enum: ComponentType, name: str, /, **kwargs) -> BaseComponent:
         """Update an existing component by type/name; never creates missing components."""
diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py
index 191a27d9..4a3dad9c 100644
--- a/tests/unit/test_job.py
+++ b/tests/unit/test_job.py
@@ -11,7 +11,7 @@ import pytest
 from reme.enumeration import ComponentEnum
 from reme.application import Application
 from reme.components.base_component import BaseComponent
-from reme.components.component_registry import ComponentRegistry
+from reme.components.component_registry import ComponentRegistry, R
 from reme.components.job.background_job import BackgroundJob
 from reme.components.job.base_job import BaseJob
 from reme.components.job.cron_job import CronJob
@@ -489,3 +489,59 @@ def test_application_start_failure_propagates_and_closes_started_components():
         assert not app._started_components
 
     asyncio.run(run())
+
+
+def test_cancelled_application_close_finishes_remaining_cleanup(tmp_path):
+    """Cancellation still closes other components and shuts down the owned pool."""
+
+    async def run():
+        closing = asyncio.Event()
+        closed = []
+
+        class ShutdownComponent(BaseComponent):
+            component_type = ComponentEnum.FILE_CHUNKER
+
+            async def _close(self):
+                try:
+                    if self.name == "b":
+                        closing.set()
+                        await asyncio.Event().wait()
+                finally:
+                    closed.append(self.name)
+
+        with R.preserve(allow_mutation=True):
+            R.register(ShutdownComponent, "test_shutdown")
+            app = Application(
+                workspace_dir=str(tmp_path),
+                enable_logo=False,
+                log_to_console=False,
+                log_to_file=False,
+                thread_pool_max_workers=1,
+                service={"backend": "cli"},
+                components={"file_chunker": {name: {"backend": "test_shutdown"} for name in ("a", "b")}},
+            )
+
+        await app.start()
+        pool = app.context.thread_pool
+        assert pool is not None
+        close_task = asyncio.create_task(app.close())
+        try:
+            await asyncio.wait_for(closing.wait(), timeout=5)
+            close_task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await close_task
+
+            assert closed == ["b", "a"]
+            assert all(not component.is_started for component in app.context.components["file_chunker"].values())
+            assert not app.is_started
+            assert app.context.thread_pool is None
+            with pytest.raises(RuntimeError, match="cannot schedule new futures after shutdown"):
+                pool.submit(lambda: None)
+        finally:
+            close_task.cancel()
+            await asyncio.gather(close_task, return_exceptions=True)
+            for component in app.context.components["file_chunker"].values():
+                await component.close()
+            pool.shutdown(wait=True)
+
+    asyncio.run(run())