AgentServer: integrate resilient tasks with memory snapshot session activation

sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md

@@ -8,6 +8,12 @@
 
 ### Bugs Fixed
 
+- Preserve unfinished task records and saved inputs when failed restore
+  activation stops workers before rolling back session configuration.
+- Keep the lease instance stable on same-restore retries, retain worker and
+  lease handles before shutdown is signalled, and defer steering promotion
+  until rollback cleanup has finished.
+
 ### Other Changes
 
 ## 2.3.0b1 (2026-10-15)
@@ -30,6 +36,13 @@
   the restored session through the ambient request context during callbacks.
   Environment-backed `AgentConfig` values are refreshed in place before
   callbacks run, and `AgentSessionContext` is marked experimental.
+- Separate process-side task registration from session-bound task activation.
+  `FOUNDRY_MEMORY_SNAPSHOT_ENABLED=true` defers enabled tasks during capture until
+  after-restore; absent or false keeps ordinary cold startup. Resilience opt-out
+  remains disabled for both initialization paths. Preserve
+  best-effort recovery, checkpoint initialization phases by restore ID, wait for
+  owned workers to stop before incompatible configuration rollback, and stop
+  managers created after startup at shutdown.
 
 ### Bugs Fixed
 

sdk/agentserver/azure-ai-agentserver-core/README.md

@@ -2,6 +2,30 @@
 
 The `azure-ai-agentserver-core` package provides the foundation host framework for building Azure AI Hosted Agent containers. It handles the protocol-agnostic infrastructure — health probes, graceful shutdown, OpenTelemetry tracing, and ASGI serving — so that protocol packages can focus on their endpoint logic.
 
+## Memory snapshot task initialization
+
+The draft integration separates task definitions from session-bound TaskManager
+initialization. Resilient tasks must still be explicitly enabled.
+
+`FOUNDRY_MEMORY_SNAPSHOT_ENABLED=true` defers enabled tasks in the temporary
+capture sandbox until the platform after-restore hook, after session hydration
+and dependency callbacks. An absent value or `false` activates enabled tasks
+during ordinary cold/local startup. `set_resilient_tasks_enabled(False)` keeps
+the subsystem disabled at startup and after-restore regardless of this variable.
+Present values must be literal `true` or `false` (case-insensitive). The variable
+does not gate running tasks or lease renewal after activation.
+
+Task startup checkpoints completed phases by restore ID and owns the manager
+through shutdown, including managers created after lifespan startup. Existing
+best-effort recovery dispatches reclaimed work immediately during the hook.
+FHA admits new requests after successful restore.
+
+Author startup task submissions are not queued or replayed by this integration.
+No session-start callback or launch-intent registry is included. Local tests are
+not proof of hosted golden restoration. See the [design specification](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/agentserver/azure-ai-agentserver-core/docs/memory-snapshot-task-integration.md)
+for lifecycle ordering, launcher compatibility requirements, validation, and
+the unresolved startup-submission options.
+
 ## Getting started
 
 ### Install the package

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py

@@ -15,6 +15,7 @@
 import os
 import traceback
 from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
 from typing import Any, TypeVar
 
 from .._config import AgentConfig
@@ -297,11 +298,11 @@ def get_task_manager() -> TaskManager:
         )
 
         raise TaskManagerNotInitialized(
-            "TaskManager not initialized: the resilient task subsystem is not enabled. "
+            "TaskManager not initialized: resilient tasks are disabled or session initialization is incomplete. "
             "Durable tasks and crash recovery are opt-in - call "
             "set_resilient_tasks_enabled(True) before host startup (e.g. at import time), "
             "or, for the responses protocol, construct the host with "
-            "ResponsesServerOptions(resilient_background=True)."  # pylint: disable=implicit-str-concat
+            "ResponsesServerOptions(resilient_background=True), and wait for session activation."
         )
     return _manager
 
@@ -393,6 +394,16 @@ def __init__(
         self._pending_input_count: int = 0
 
 
+@dataclass
+class _TaskInitializationPhase:
+    completed: bool = False
+    completed_restore_id: str | None = None
+
+    def complete(self, restore_id: str | None) -> None:
+        self.completed_restore_id = restore_id
+        self.completed = True
+
+
 class TaskManager:  # pylint: disable=too-many-instance-attributes,protected-access,broad-exception-caught
     """Lifecycle orchestrator for resilient tasks.
 
@@ -430,12 +441,17 @@ def __init__(
         self._instance_id = generate_instance_id()
         self._materialization_restore_id: str | None = None
         self._shutdown_event = shutdown_event or asyncio.Event()
+        self._preserve_recoverable_tasks_on_shutdown = False
         self._shutdown_grace_seconds = shutdown_grace_seconds
         self._active_generation_future: dict[str, asyncio.Future[Any]] = {}
         self._pending_steering_futures: dict[str, list[asyncio.Future[Any]]] = {}
         #   Layer 2: periodic recovery scan task. Created
         # at startup() time; cancelled at shutdown().
         self._periodic_recovery_task: asyncio.Task[None] | None = None
+        self._startup_lock = asyncio.Lock()
+        self._startup_phases = {
+            name: _TaskInitializationPhase() for name in ("construct", "register", "recovery", "periodic", "ready")
+        }
         #   / C-WQ-1..3 — per-task write-queue
         # registry. A single asyncio.Lock per task_id serializes all
         # in-process PATCHes against that task so etag conflicts become
@@ -740,40 +756,61 @@ async def _cancel_queued_steering_input(  # pylint: disable=unused-argument
         if not future.done():
             future.set_exception(TaskCancelled())
 
-    async def startup(self) -> None:
+    @property
+    def _has_live_workers(self) -> bool:
+        return (
+            bool(self._active_tasks)
+            or (self._periodic_recovery_task is not None and not self._periodic_recovery_task.done())
+            or any(not worker.done() for worker in self._timeout_watchdogs.values())
+        )
+
+    async def startup(self, *, restore_id: str | None = None) -> None:
         """Initialize the manager and recover stale tasks.
 
-        Called by ``AgentServerHost`` during lifespan startup.
-        """
-        logger.info(
-            "TaskManager starting (owner=%s, instance=%s, hosted=%s)",
-            self._lease_owner,
-            self._instance_id,
-            self._config.is_hosted,
-        )
-        # Pick up descriptors registered at import time (for recovery)
-        from ._decorator import (  # pylint: disable=import-outside-toplevel
-            _REGISTERED_DESCRIPTORS,
-        )
+        Called during ordinary session startup or after memory restoration.
 
-        for fn_name, fn, opts in _REGISTERED_DESCRIPTORS:
-            self._resume_callbacks[fn_name] = fn
-            self._resume_opts[fn_name] = opts
+        :keyword restore_id: Platform restore identity, or None for cold startup.
+        :paramtype restore_id: str or None
+        """
+        async with self._startup_lock:
+            if self._shutdown_event.is_set():
+                raise RuntimeError("Cannot initialize a stopped TaskManager")
+            logger.info(
+                "TaskManager starting (owner=%s, instance=%s, hosted=%s, restore_id=%s)",
+                self._lease_owner,
+                self._instance_id,
+                self._config.is_hosted,
+                restore_id,
+            )
+            self._startup_phases["construct"].complete(restore_id)
+            register = self._startup_phases["register"]
+            if not register.completed:
+                from ._decorator import (  # pylint: disable=import-outside-toplevel
+                    _REGISTERED_DESCRIPTORS,
+                )
 
-        await self._recover_stale_tasks()
+                for fn_name, fn, opts in _REGISTERED_DESCRIPTORS:
+                    self._resume_callbacks[fn_name] = fn
+                    self._resume_opts[fn_name] = opts
+            register.complete(restore_id)
 
-        #   Layer 2: start the periodic recovery task.
-        # Reads _PERIODIC_RECOVERY_INTERVAL_SECONDS at spawn time;
-        # tests monkey-patch the constant to drive the scan
-        # deterministically.
-        try:
-            loop = asyncio.get_running_loop()
-            self._periodic_recovery_task = loop.create_task(self._periodic_recovery_loop())
-        except RuntimeError:
-            # No running loop (called from outside async context); skip
-            # — the layer-1 startup scan above still covered the
-            # initial reclaim pass.
-            pass
+            recovery = self._startup_phases["recovery"]
+            if not recovery.completed or recovery.completed_restore_id != restore_id:
+                await self._recover_stale_tasks()
+                recovery.complete(restore_id)
+
+            scanner = self._periodic_recovery_task
+            if scanner is not None and scanner.done():
+                self._startup_phases["periodic"].completed = False
+                self._startup_phases["ready"].completed = False
+                self._periodic_recovery_task = None
+                if not scanner.cancelled():
+                    scanner.result()
+                logger.warning("Restarting stopped periodic task recovery scanner")
+            if self._periodic_recovery_task is None:
+                self._periodic_recovery_task = asyncio.create_task(self._periodic_recovery_loop())
+            self._startup_phases["periodic"].complete(restore_id)
+            self._startup_phases["ready"].complete(restore_id)
 
     async def _periodic_recovery_loop(self) -> None:
         """Layer 2: periodic background recovery scan.
@@ -802,12 +839,61 @@ async def _periodic_recovery_loop(self) -> None:
             except Exception:  # pylint: disable=broad-exception-caught
                 logger.warning("Periodic recovery scan iteration failed", exc_info=True)
 
-    async def shutdown(self) -> None:
+    async def _cancel_workers(self, *, wait: bool, retained_workers: set[asyncio.Task[Any]] | None = None) -> None:
+        """Cancel owned execution, renewal, and watchdog tasks.
+
+        :keyword wait: Whether to wait for the cancelled tasks to finish.
+        :paramtype wait: bool
+        :keyword retained_workers: Workers retained before shutdown was signalled.
+        :paramtype retained_workers: set[asyncio.Task[Any]] | None
+        """
+        workers = set(retained_workers or ())
+        for active in list(self._active_tasks.values()):
+            active.renewal_cancel.set()
+            if active.renewal_task and not active.renewal_task.done():
+                workers.add(active.renewal_task)
+            if not active.execution_task.done():
+                workers.add(active.execution_task)
+        for watchdog in self._timeout_watchdogs.values():
+            if not watchdog.done():
+                workers.add(watchdog)
+        for worker in workers:
+            if not worker.done():
+                worker.cancel()
+        if not (wait and workers):
+            return
+        results = await asyncio.gather(*workers, return_exceptions=True)
+        for result in results:
+            if isinstance(result, Exception):
+                logger.warning(
+                    "Task worker failed during shutdown",
+                    exc_info=(type(result), result, result.__traceback__),
+                )
+
+    async def shutdown(self, *, wait_for_workers: bool = False, preserve_recoverable_tasks: bool = False) -> None:
         """Signal shutdown on all active tasks and force-expire leases.
 
         Called by ``AgentServerHost`` during lifespan shutdown.
+
+        :keyword wait_for_workers: Also wait for the cancelled execution, renewal,
+            and watchdog tasks to finish. Used before session configuration is
+            rolled back so handlers cannot observe the later binding.
+        :paramtype wait_for_workers: bool
+        :keyword preserve_recoverable_tasks: Stop workers before expiring leases
+            without treating rollback interruption as terminal cancellation.
+            Always waits for workers; used only for failed restore rollback.
+        :paramtype preserve_recoverable_tasks: bool
         """
         logger.info("TaskManager shutting down")
+        self._preserve_recoverable_tasks_on_shutdown = preserve_recoverable_tasks
+        retained_active = tuple(self._active_tasks.values()) if preserve_recoverable_tasks else ()
+        retained_workers: set[asyncio.Task[Any]] = (
+            set(self._timeout_watchdogs.values()) if preserve_recoverable_tasks else set()
+        )
+        for active in retained_active:
+            retained_workers.add(active.execution_task)
+            if active.renewal_task is not None:
+                retained_workers.add(active.renewal_task)
         self._shutdown_event.set()
 
         #   Layer 2: stop the periodic recovery scan task.
@@ -864,9 +950,15 @@ async def shutdown(self) -> None:
             except asyncio.CancelledError:
                 logger.info("TaskManager shutdown grace period interrupted")
 
+        remaining = list(self._active_tasks.values())
+        if preserve_recoverable_tasks:
+            remaining = list({id(active): active for active in (*retained_active, *remaining)}.values())
+            # Finish old-binding writes before making the records reclaimable.
+            await self._cancel_workers(wait=True, retained_workers=retained_workers)
+
         # Force-expire all leases. Tolerate cancellation here too.
         try:
-            for active in list(self._active_tasks.values()):
+            for active in remaining:
                 try:
                     await self._provider.update(
                         active.task_id,
@@ -887,15 +979,16 @@ async def shutdown(self) -> None:
 
         # Cancel all renewal and execution tasks. Always do this so handlers
         # listening on the cancellation signal wake up and exit cleanly.
-        for active in self._active_tasks.values():
-            active.renewal_cancel.set()
-            if active.renewal_task and not active.renewal_task.done():
-                active.renewal_task.cancel()
-            if not active.execution_task.done():
-                active.execution_task.cancel()
+        if not preserve_recoverable_tasks:
+            await self._cancel_workers(wait=wait_for_workers)
 
         self._active_tasks.clear()
-        set_task_manager(None)
+        self._timeout_watchdogs.clear()
+        for phase in self._startup_phases.values():
+            phase.completed = False
+            phase.completed_restore_id = None
+        if _manager is self:
+            set_task_manager(None)
 
     async def create_and_run(
         self,
@@ -2087,6 +2180,14 @@ async def _execute_task_loop(  # pylint: disable=too-many-statements,too-many-br
 
             except asyncio.CancelledError:
                 renewal_cancel.set()
+                if self._preserve_recoverable_tasks_on_shutdown:
+                    from ._exceptions import (  # pylint: disable=import-outside-toplevel
+                        TaskDeferred,
+                    )
+
+                    if not current_result_future.done():
+                        current_result_future.set_exception(TaskDeferred())
+                    break
                 # asyncio.CancelledError is the cooperative-cancel path —
                 # the handler chose to raise it (or the framework signalled
                 # cancel via ctx.cancel and the handler did not catch).
@@ -2350,6 +2451,9 @@ async def _try_drain_steering(  # pylint: disable=too-many-branches,too-many-sta
         # (detect-not-clobber) so a genuine cross-process write is detected;
         # cross-process conflicts retry OUTSIDE the lock via the recursion
         # below (the per-task ``asyncio.Lock`` is non-reentrant).
+        if self._preserve_recoverable_tasks_on_shutdown:
+            return None
+
         drain_conflict: BaseException | None = None
         async with self._get_task_write_lock(task_id):
             task_info = await self._provider_get_tracked(task_id)