optio-core 0.2.2__tar.gz → 0.2.4__tar.gz

{optio_core-0.2.2/src/optio_core.egg-info → optio_core-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optio-core
-Version: 0.2.2
+Version: 0.2.4
 Summary: Reusable async process management library
 Author-email: Kristof Csillag <kristof.csillag@deai-labs.com>
 License-Expression: Apache-2.0

{optio_core-0.2.2 → optio_core-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "optio-core"
-version = "0.2.2"
+version = "0.2.4"
 description = "Reusable async process management library"
 readme = "README.md"
 license = "Apache-2.0"

{optio_core-0.2.2 → optio_core-0.2.4}/src/optio_core/__init__.py

@@ -34,7 +34,7 @@ __all__ = [
     "get_process", "list_processes",
     "block_launches", "unblock_launches",
     "group_cancel", "group_cancel_and_wait",
-    "rpc_server", "store", "MongoStore",
+    "rpc_server", "mongo_store", "MongoStore",
 ]
 
 
@@ -47,7 +47,8 @@ def __getattr__(name: str):
     """
     if name == "rpc_server":
         return _instance.rpc_server
-    if name == "store":
+    if name == "mongo_store":
         # Runtime-populated: the (db, prefix) binding only exists after init().
-        return _instance.store
+        # NB: not "store" — that name is the optio_core.store submodule.
+        return _instance.mongo_store
     raise AttributeError(f"module 'optio_core' has no attribute {name!r}")

{optio_core-0.2.2 → optio_core-0.2.4}/src/optio_core/lifecycle.py

@@ -82,20 +82,24 @@ class Optio:
         self._owned_rpc_server: bool = False
         self._scheduler: ProcessScheduler | None = None
         self._running = False
+        self._shutting_down = False
+        self._shutdown_task: "asyncio.Task | None" = None
         self._heartbeat_task: asyncio.Task | None = None
         self._supervisor_task: asyncio.Task | None = None
         self._launch_blocks: dict[uuid.UUID, _BlockEntry] = {}
 
     @property
-    def store(self) -> MongoStore:
+    def mongo_store(self) -> MongoStore:
         """The (db, prefix) binding this instance was initialized with.
 
+        Named ``mongo_store`` (not ``store``) because ``optio_core.store`` is a
+        submodule — a module-level ``store`` attribute would be shadowed by it.
         Hand it as-is to helpers that need db+prefix (e.g.
         ``optio_claudecode.delete_seed``) instead of re-deriving the prefix
         from your own config. Raises if accessed before ``init()``.
         """
         if self._config is None:
-            raise RuntimeError("optio.store accessed before init()")
+            raise RuntimeError("optio.mongo_store accessed before init()")
         return MongoStore(db=self._config.mongo_db, prefix=self._config.prefix)
 
     async def init(
@@ -389,6 +393,11 @@ class Optio:
         background and the outcome is ok=True with `proc` populated from a
         post-schedule re-read (state will typically be 'scheduled' or
         already advanced to 'running' depending on event-loop timing)."""
+        # Stop sign: once shutdown has begun, refuse new launches so nothing
+        # half-boots into a shutting-down engine. Covers scheduler, RPC, and
+        # programmatic launch paths.
+        if self._shutting_down:
+            return LaunchOutcome(ok=False, reason="shutting-down")
         proc = await self._resolve(process_id)
         if proc is None:
             return LaunchOutcome(ok=False, reason="not-found")
@@ -877,10 +886,15 @@ class Optio:
         # Set up signal handlers (only works in main thread)
         try:
             loop = asyncio.get_event_loop()
+
+            def _request_shutdown() -> None:
+                # Store the task so run() can await it to completion before the
+                # loop is torn down; guard against a second signal.
+                if self._shutdown_task is None:
+                    self._shutdown_task = asyncio.create_task(self.shutdown())
+
             for sig in (signal.SIGTERM, signal.SIGINT):
-                loop.add_signal_handler(
-                    sig, lambda: asyncio.create_task(self.shutdown()),
-                )
+                loop.add_signal_handler(sig, _request_shutdown)
         except (NotImplementedError, RuntimeError):
             pass  # Signal handlers not available (e.g., in tests)
 
@@ -906,6 +920,16 @@ class Optio:
         try:
             await self._shutdown_event.wait()
         finally:
+            # A signal dispatches shutdown() as a detached task. Wait for its
+            # cooperative drain to finish before returning — otherwise main()
+            # returns and asyncio.run() tear-down cancels the drain mid-capture.
+            # (When shutdown() is called directly rather than via signal,
+            # _shutdown_task is None and the caller already awaited it.)
+            if self._shutdown_task is not None:
+                try:
+                    await self._shutdown_task
+                except Exception:
+                    logger.exception("shutdown task raised")
             await self._scheduler.stop()
 
     async def shutdown(self, grace_seconds: float | None = None) -> None:
@@ -920,67 +944,75 @@ class Optio:
         """
         logger.info("Shutdown requested")
         self._running = False
+        # Stop sign: refuse new launches (scheduler, RPC, programmatic) for the
+        # entire drain so nothing half-boots into a shutting-down engine.
+        self._shutting_down = True
 
-        # 1. Heartbeat
-        if self._heartbeat_task:
-            self._heartbeat_task.cancel()
-            try:
-                await self._heartbeat_task
-            except asyncio.CancelledError:
-                pass
-            self._heartbeat_task = None
-
-        # 2. Signal shutdown to the main run() loop
-        if hasattr(self, '_shutdown_event'):
-            self._shutdown_event.set()
+        try:
+            # 1. Heartbeat
+            if self._heartbeat_task:
+                self._heartbeat_task.cancel()
+                try:
+                    await self._heartbeat_task
+                except asyncio.CancelledError:
+                    pass
+                self._heartbeat_task = None
+
+            # 2. RPC server (clamator)
+            if self.rpc_server is not None and self._owned_rpc_server:
+                grace_ms = int(
+                    (grace_seconds if grace_seconds is not None else self._config.cancel_grace_seconds) * 1000
+                )
+                await self.rpc_server.stop(grace_ms=grace_ms)
 
-        # 2b. RPC server (clamator)
-        if self.rpc_server is not None and self._owned_rpc_server:
-            grace_ms = int(
-                (grace_seconds if grace_seconds is not None else self._config.cancel_grace_seconds) * 1000
+            # 3. Cancel everything via the unified mechanism.
+            grace = (
+                grace_seconds
+                if grace_seconds is not None
+                else self._config.cancel_grace_seconds
             )
-            await self.rpc_server.stop(grace_ms=grace_ms)
-
-        # 3. Cancel everything via the unified mechanism.
-        grace = (
-            grace_seconds
-            if grace_seconds is not None
-            else self._config.cancel_grace_seconds
-        )
-        if self._executor:
-            now_mono = time.monotonic()
-            for oid in list(self._executor._cancellation_flags.keys()):
-                entry = self._executor._cancellation_flags.get(oid)
-                if entry is None:
-                    continue
-                entry.flag.set()
-                if entry.deadline is None:
-                    entry.deadline = now_mono + grace
-
-            # Wait for entries to drain. The supervisor handles force-cancel.
-            ceiling = time.monotonic() + grace + 5.0
-            while self._executor._cancellation_flags and time.monotonic() < ceiling:
-                await asyncio.sleep(0.1)
-
-            # Belt and braces: anything still left, force-cancel directly.
-            # (Handles the case where the supervisor was slow or already stopped.)
-            for oid in list(self._executor._cancellation_flags.keys()):
-                await self._executor.force_cancel(oid)
-
-        # 4. Stop supervisor (after final force-cancel pass).
-        if self._supervisor_task:
-            self._supervisor_task.cancel()
-            try:
-                await self._supervisor_task
-            except asyncio.CancelledError:
-                pass
-            self._supervisor_task = None
-
-        # 5. Redis
-        if self._redis:
-            await self._redis.aclose()
-
-        logger.info("Shutdown complete")
+            if self._executor:
+                now_mono = time.monotonic()
+                for oid in list(self._executor._cancellation_flags.keys()):
+                    entry = self._executor._cancellation_flags.get(oid)
+                    if entry is None:
+                        continue
+                    entry.flag.set()
+                    if entry.deadline is None:
+                        entry.deadline = now_mono + grace
+
+                # Wait for entries to drain. The supervisor handles force-cancel.
+                ceiling = time.monotonic() + grace + 5.0
+                while self._executor._cancellation_flags and time.monotonic() < ceiling:
+                    await asyncio.sleep(0.1)
+
+                # Belt and braces: anything still left, force-cancel directly.
+                # (Handles the case where the supervisor was slow or already stopped.)
+                for oid in list(self._executor._cancellation_flags.keys()):
+                    await self._executor.force_cancel(oid)
+
+            # 4. Stop supervisor (after final force-cancel pass).
+            if self._supervisor_task:
+                self._supervisor_task.cancel()
+                try:
+                    await self._supervisor_task
+                except asyncio.CancelledError:
+                    pass
+                self._supervisor_task = None
+
+            # 5. Redis
+            if self._redis:
+                await self._redis.aclose()
+
+            logger.info("Shutdown complete")
+        finally:
+            # Unblock run() ONLY after the drain has fully run (or failed). The
+            # finally guarantees run() still unblocks if the drain raised —
+            # otherwise run() would hang until the external SIGKILL backstop.
+            # (Moved here from an early "step 2": setting it before the drain let
+            # run() return and asyncio.run() tear-down cancel the drain.)
+            if hasattr(self, '_shutdown_event'):
+                self._shutdown_event.set()
 
     async def _reconcile_interrupted_processes(self) -> None:
         """Mark processes left in active states by a previous session as failed.

{optio_core-0.2.2 → optio_core-0.2.4}/src/optio_core/models.py

@@ -15,6 +15,7 @@ class LaunchOutcome:
     ok: bool
     reason: Literal[
         "not-found", "not-launchable", "launch-blocked", "no-resume-support",
+        "shutting-down",
     ] | None = None
     proc: dict[str, Any] | None = None
 

{optio_core-0.2.2 → optio_core-0.2.4/src/optio_core.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optio-core
-Version: 0.2.2
+Version: 0.2.4
 Summary: Reusable async process management library
 Author-email: Kristof Csillag <kristof.csillag@deai-labs.com>
 License-Expression: Apache-2.0

{optio_core-0.2.2 → optio_core-0.2.4}/src/optio_core.egg-info/SOURCES.txt

@@ -55,6 +55,7 @@ tests/test_progress_throttle.py
 tests/test_resync.py
 tests/test_resync_cancel_stale.py
 tests/test_scheduler.py
+tests/test_shutdown_drain_completion.py
 tests/test_state_machine.py
 tests/test_store.py
 tests/test_task_ttl.py

{optio_core-0.2.2 → optio_core-0.2.4}/tests/test_cancel_propagation.py

@@ -578,8 +578,19 @@ async def test_force_cancel_cascade_catches_late_optout_child(mongo_db):
     await late_child_spawned.wait()
 
     # Within force-cancel window, late child exists in DB and is active.
-    late_proc = await get_process_by_process_id(mongo_db, prefix, "late")
-    assert late_proc is not None
+    # spawn() schedules child persistence as a background task
+    # (context.py: asyncio.create_task(_run())); it does not await the DB
+    # write. Poll for the late child to appear and be active within the
+    # force-cancel window instead of reading eagerly (the eager read raced
+    # cold-start persistence and failed when this test ran first/alone).
+    late_proc = None
+    end = _time.monotonic() + 0.4
+    while _time.monotonic() < end:
+        late_proc = await get_process_by_process_id(mongo_db, prefix, "late")
+        if late_proc is not None and late_proc["status"]["state"] in {"running", "scheduled"}:
+            break
+        await asyncio.sleep(0.01)
+    assert late_proc is not None, "late child never persisted/active within window"
     assert late_proc["status"]["state"] in {"running", "scheduled"}
 
     # After grace + cascade, late child reaches terminal.

optio_core-0.2.4/tests/test_shutdown_drain_completion.py

@@ -0,0 +1,150 @@
+"""Shutdown must drain cooperating tasks BEFORE unblocking run().
+
+Spec: docs/superpowers/specs/2026-06-02-graceful-engine-shutdown-capture-design.md
+(Addendum). Regression guard for the run()/shutdown() ordering race: the
+signal handler dispatched shutdown() as a detached task that set
+_shutdown_event before the cooperative drain, so run() returned and
+asyncio.run() teardown cancelled the drain mid-capture.
+"""
+import asyncio
+
+import pytest
+
+from optio_core.lifecycle import Optio
+from optio_core.models import TaskInstance
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_shutdown_event_fires_only_after_drain(mongo_db):
+    prefix = "shut1"
+    running = asyncio.Event()
+    captured = asyncio.Event()
+
+    async def task(ctx):
+        running.set()
+        # Cooperate with cancel, then simulate a snapshot capture that takes
+        # real time inside the cooperative window.
+        for _ in range(200):
+            if ctx.cancellation_flag.is_set():
+                break
+            await asyncio.sleep(0.02)
+        await asyncio.sleep(0.3)  # "capture"
+        captured.set()
+
+    ti = TaskInstance(process_id="t.cap", name="Cap", params={}, execute=task)
+
+    async def gen(_s, _f):
+        return [ti]
+
+    optio = Optio()
+    await optio.init(
+        mongo_db=mongo_db, prefix=prefix,
+        get_task_definitions=gen, cancel_grace_seconds=2.0,
+    )
+    run_task = asyncio.create_task(optio.run())
+    try:
+        await optio.launch("t.cap", session_id=None)
+        await asyncio.wait_for(running.wait(), timeout=5)
+
+        # Mimic the signal handler: dispatch shutdown as a detached task.
+        shutdown_task = asyncio.create_task(optio.shutdown())
+
+        # The invariant: by the time run()'s unblock event is set, the
+        # cooperative drain must already have completed (task captured).
+        await asyncio.wait_for(optio._shutdown_event.wait(), timeout=5)
+        assert captured.is_set(), (
+            "_shutdown_event fired before the running task finished its "
+            "cooperative capture — run() would return and the loop tear-down "
+            "would abort the drain"
+        )
+
+        await asyncio.wait_for(shutdown_task, timeout=5)
+    finally:
+        await asyncio.wait_for(run_task, timeout=5)
+
+
+async def test_launch_refused_during_shutdown(mongo_db):
+    prefix = "shut2"
+    running = asyncio.Event()
+
+    async def slow(ctx):
+        running.set()
+        for _ in range(200):
+            if ctx.cancellation_flag.is_set():
+                return
+            await asyncio.sleep(0.02)
+
+    a = TaskInstance(process_id="t.a", name="A", params={}, execute=slow)
+    b = TaskInstance(process_id="t.b", name="B", params={}, execute=slow)
+
+    async def gen(_s, _f):
+        return [a, b]
+
+    optio = Optio()
+    await optio.init(
+        mongo_db=mongo_db, prefix=prefix,
+        get_task_definitions=gen, cancel_grace_seconds=2.0,
+    )
+    run_task = asyncio.create_task(optio.run())
+    try:
+        await optio.launch("t.a", session_id=None)
+        await asyncio.wait_for(running.wait(), timeout=5)
+
+        shutdown_task = asyncio.create_task(optio.shutdown())
+        # Give shutdown() a tick to set _shutting_down at its start.
+        await asyncio.sleep(0.05)
+
+        outcome = await optio.launch("t.b", session_id=None)
+        assert outcome.ok is False
+        assert outcome.reason == "shutting-down"
+
+        await asyncio.wait_for(shutdown_task, timeout=5)
+    finally:
+        await asyncio.wait_for(run_task, timeout=5)
+
+
+async def test_shutdown_drains_parent_and_child(mongo_db):
+    prefix = "shut3"
+    parent_running = asyncio.Event()
+    child_running = asyncio.Event()
+    child_captured = asyncio.Event()
+
+    async def child(ctx):
+        child_running.set()
+        for _ in range(200):
+            if ctx.cancellation_flag.is_set():
+                break
+            await asyncio.sleep(0.02)
+        await asyncio.sleep(0.2)
+        child_captured.set()
+
+    async def parent(ctx):
+        parent_running.set()
+        await ctx.run_child(execute=child, process_id="p.child", name="Child", params={})
+
+    parent_ti = TaskInstance(process_id="p.parent", name="Parent", params={}, execute=parent)
+
+    async def gen(_s, _f):
+        return [parent_ti]
+
+    optio = Optio()
+    await optio.init(
+        mongo_db=mongo_db, prefix=prefix,
+        get_task_definitions=gen, cancel_grace_seconds=2.0,
+    )
+    run_task = asyncio.create_task(optio.run())
+    try:
+        await optio.launch("p.parent", session_id=None)
+        await asyncio.wait_for(parent_running.wait(), timeout=5)
+        await asyncio.wait_for(child_running.wait(), timeout=5)
+
+        shutdown_task = asyncio.create_task(optio.shutdown())
+        await asyncio.wait_for(optio._shutdown_event.wait(), timeout=6)
+
+        # The child cooperated and captured before the event fired.
+        assert child_captured.is_set()
+
+        await asyncio.wait_for(shutdown_task, timeout=5)
+    finally:
+        await asyncio.wait_for(run_task, timeout=5)