arkitekt-runtime 7.2.0__tar.gz → 7.3.0__tar.gz

{arkitekt_runtime-7.2.0 → arkitekt_runtime-7.3.0}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.5
 Name: arkitekt-runtime
-Version: 7.2.0
+Version: 7.3.0
 Summary: the execution core of an Arkitekt app: agent, actors, tasks, journal, state and structure serialization, shared by every runtime
 Author-email: jhnnsrs <jhnnsrs@gmail.com>
 License-Expression: MIT
 Requires-Python: >=3.11
-Requires-Dist: arkitekt-spec>=6.2
+Requires-Dist: arkitekt-spec>=6.3.1
 Requires-Dist: janus>=2.0.0
 Requires-Dist: jsonpatch>=1.33
 Requires-Dist: koil>=3.3.6

{arkitekt_runtime-7.2.0 → arkitekt_runtime-7.3.0}/arkitekt_runtime/agents/base.py

@@ -318,6 +318,8 @@ class BaseAgent(KoiledModel):
         default_factory=lambda: {}
     )
     _ran_startup_hooks: bool = PrivateAttr(default=False)
+    _entered_lifecycles: list[Any] = PrivateAttr(default_factory=list)
+    """The lifecycle hooks started in this run, in order: what follows their yield is owed."""
     # A stopped run reaches teardown from several places (the loop, leaving the
     # agent); it runs once per connect.
     _torn_down: bool = PrivateAttr(default=False)
@@ -1824,6 +1826,8 @@ class BaseAgent(KoiledModel):
         for key, hook in self._collected_startup_hooks.items():
             try:
                 previous = await self._aprevious_state_kwargs(hook)
+                if hasattr(hook, "afinish"):
+                    self._entered_lifecycles.append(hook)
                 answer = await asyncio.wait_for(
                     hook.arun(
                         app_context=app_context,
@@ -1887,6 +1891,25 @@ class BaseAgent(KoiledModel):
                 hook_error.__cause__ = e
                 logger.error(hook_error, exc_info=hook_error)
 
+        # After the shutdown hooks: they may take what these set up.
+        await self._afinish_lifecycles()
+
+    async def _afinish_lifecycles(self) -> None:
+        """Run what follows each lifecycle hook's yield, the last one entered first.
+
+        Like a shutdown hook, one that fails or times out is logged, and the rest still run.
+        """
+        from arkitekt_spec.declare.agents.hooks.errors import ShutdownHookError
+
+        entered, self._entered_lifecycles = self._entered_lifecycles, []
+        for hook in reversed(entered):
+            try:
+                await asyncio.wait_for(hook.afinish(), timeout=self.shutdown_hook_timeout)
+            except Exception as e:
+                hook_error = ShutdownHookError(f"Lifecycle hook {hook.func.__name__} failed")
+                hook_error.__cause__ = e
+                logger.error(hook_error, exc_info=hook_error)
+
     @property
     def registered_agent_id(self) -> str | None:
         """The id this agent's backend assigned it, once registered.
@@ -1928,7 +1951,13 @@ class BaseAgent(KoiledModel):
         """
         # A state is built unrestricted by its startup hook; adopting it (in
         # ainit_states) is what makes later changes need their locks.
-        hook_return = await self.arun_startup_hooks(app_context=app_context)
+        try:
+            hook_return = await self.arun_startup_hooks(app_context=app_context)
+        except BaseException:
+            # No shutdown hook runs for a start that failed here, but the lifecycle
+            # hooks before the failing one are entered.
+            await self._afinish_lifecycles()
+            raise
         # From here on the app has set up resources, so teardown owes it the
         # shutdown hooks (even if the rest of the startup fails).
         self._ran_startup_hooks = True

{arkitekt_runtime-7.2.0 → arkitekt_runtime-7.3.0}/arkitekt_runtime/agents/hooks/materialize.py

@@ -4,6 +4,7 @@ from typing import Any
 
 from arkitekt_spec.declare.agents.hooks.declare import (
     BackgroundDeclaration,
+    LifecycleHookDeclaration,
     ShutdownHookDeclaration,
     StartupHookDeclaration,
 )
@@ -12,7 +13,11 @@ from arkitekt_runtime.agents.hooks.background import (
     WrappedThreadedBackgroundTask,
 )
 from arkitekt_runtime.agents.hooks.shutdown import ThreadedShutdownHook, WrappedShutdownHook
-from arkitekt_runtime.agents.hooks.startup import ThreadedStartupHook, WrappedStartupHook
+from arkitekt_runtime.agents.hooks.startup import (
+    ThreadedStartupHook,
+    WrappedLifecycleHook,
+    WrappedStartupHook,
+)
 
 
 def materialize_hook(declaration: Any) -> Any:  # noqa: ANN401
@@ -24,7 +29,9 @@ def materialize_hook(declaration: Any) -> Any:  # noqa: ANN401
     if hasattr(declaration, "arun"):
         return declaration
     func, registry = declaration.func, declaration.structure_registry
-    if isinstance(declaration, StartupHookDeclaration):
+    if isinstance(declaration, LifecycleHookDeclaration):
+        cls = WrappedLifecycleHook
+    elif isinstance(declaration, StartupHookDeclaration):
         cls = ThreadedStartupHook if declaration.threaded else WrappedStartupHook
     elif isinstance(declaration, BackgroundDeclaration):
         cls = WrappedThreadedBackgroundTask if declaration.threaded else WrappedBackgroundTask

{arkitekt_runtime-7.2.0 → arkitekt_runtime-7.3.0}/arkitekt_runtime/agents/hooks/startup.py

@@ -1,16 +1,20 @@
 """Hooks for the agent"""
 
+import asyncio
 from typing import (
     TYPE_CHECKING,
     Any,
 )
 from arkitekt_spec.declare.agents.types import BoundApp
 from koil.bridge import run_threaded
-from arkitekt_spec.declare.agents.hooks.errors import StartupHookError
+from arkitekt_spec.declare.agents.hooks.errors import ShutdownHookError, StartupHookError
 from arkitekt_spec.declare.agents.hooks.registry import (
     StartupHookReturns,
 )
-from arkitekt_spec.declare.agents.hooks.declare import StartupWithVariables
+from arkitekt_spec.declare.agents.hooks.declare import (
+    LifecycleWithVariables,
+    StartupWithVariables,
+)
 
 if TYPE_CHECKING:
     from arkitekt_spec.declare.structures.registry import StructureRegistry
@@ -118,4 +122,77 @@ class ThreadedStartupHook(StartupWithVariables):
         return _split_returns(self, ensure_return_as_tuple(parsed_returns))
 
 
+class WrappedLifecycleHook(LifecycleWithVariables):
+    """A lifecycle hook, driven by a task of its own.
+
+    The agent starts and tears down in different tasks, and what an async generator enters
+    (a cancel scope, ``asyncio.timeout``) must be left in the task that entered it.
+    """
+
+    def __init__(
+        self, func: Any, structure_registry: "StructureRegistry | None" = None  # noqa: ANN401
+    ) -> None:
+        super().__init__(func, structure_registry)
+        self._task: asyncio.Task[None] | None = None
+        self._stop = asyncio.Event()
+
+    async def arun(
+        self,
+        app_context: Any,  # noqa: ANN401
+        bound_app: BoundApp | None = None,
+        previous_states: dict[str, Any] | None = None,
+    ) -> StartupHookReturns:
+        """Run the hook up to its yield, in its own task."""
+        kwargs = self.get_startup_kwargs(app_context, bound_app, previous_states)
+        entered: asyncio.Future[Any] = asyncio.get_running_loop().create_future()
+        self._stop = asyncio.Event()
+        self._task = asyncio.create_task(self._adrive(kwargs, entered))
+        try:
+            yielded = await entered
+        except BaseException:
+            # Its setup failed, or ran out of time: there is nothing left to finish.
+            task, self._task = self._task, None
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+            raise
+        return _split_returns(self, ensure_return_as_tuple(yielded))
+
+    async def afinish(self) -> None:
+        """Run what follows the yield. Nothing if the hook is not entered."""
+        task, self._task = self._task, None
+        if task is None:
+            return
+        self._stop.set()
+        await task
+
+    async def _adrive(self, kwargs: dict[str, Any], entered: "asyncio.Future[Any]") -> None:
+        name = self.func.__name__
+        generator = self.func(**kwargs)
+        try:
+            try:
+                yielded = await anext(generator)
+            except StopAsyncIteration:
+                if not entered.done():
+                    entered.set_exception(StartupHookError(f"Lifecycle hook {name} did not yield"))
+                return
+            except asyncio.CancelledError:
+                entered.cancel()
+                raise
+            except Exception as e:
+                if not entered.done():
+                    entered.set_exception(e)
+                return
+            if entered.done():
+                return
+            entered.set_result(yielded)
+            await self._stop.wait()
+            try:
+                await anext(generator)
+            except StopAsyncIteration:
+                return
+            raise ShutdownHookError(f"Lifecycle hook {name} yielded more than once")
+        finally:
+            await generator.aclose()
+
+
 # --- Implementation ---

{arkitekt_runtime-7.2.0 → arkitekt_runtime-7.3.0}/pyproject.toml

@@ -10,7 +10,7 @@ license = "MIT"
 # server) executes actions the same way, so this is what they share. It carries no
 # transport and no client: no websockets, no rath, no graphql-core.
 dependencies = [
-    "arkitekt-spec>=6.2",
+    "arkitekt-spec>=6.3.1",
     "pydantic>=2.11",
     "koil>=3.3.6",
     "janus>=2.0.0",