From 889dab714758c1e307289dad9224abc93662c2da Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Thu, 30 Jul 2026 19:52:05 +0530 Subject: [PATCH 1/7] [agentserver-core] Gate resilient TaskManager auto-init behind durable-task opt-in AgentServerHost's lifespan unconditionally initialized the resilient TaskManager at startup, issuing a blocking hosted task-store list() round-trip (plus DefaultAzureCredential token acquisition) BEFORE the lifespan yield. This gated server readiness on a network call even for apps that never use the task primitive (e.g. invocations-only hosts), adding tens of seconds of first-connection latency with nothing to recover. Gate auto-init on a non-empty _REGISTERED_DESCRIPTORS registry, which every durable decorator populates: @task and @multi_turn_task both construct a Task whose __init__ registers (MultiTurnTask wraps an inner Task, so it registers too). Declaring a durable task is now the opt-in for standing up the TaskManager and its recovery scan; apps that never opt in skip it entirely. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5222ab61-7d5e-4642-b75f-ef365e375f4f --- .../azure-ai-agentserver-core/CHANGELOG.md | 1 + .../azure/ai/agentserver/core/_base.py | 76 +++++-- .../tests/tasks/test_task_manager_optin.py | 193 ++++++++++++++++++ 3 files changed, 253 insertions(+), 17 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 50110d7b7c01..591d249b7d79 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -8,6 +8,7 @@ ### Bugs Fixed +- `AgentServerHost` now only initializes the resilient `TaskManager` at startup when the app has declared a durable task via `@task` / `@multi_turn_task`. Plain servers (e.g. invocations-only hosts) that never use the task primitive no longer pay the task-store startup cost — a blocking hosted task-store `list()` round-trip plus credential-token acquisition that could gate server readiness (and add tens of seconds of connect latency) while having nothing to recover. - Cleaned up `AgentServerHost` public signatures so inherited middleware typing does not expose Starlette private type aliases. ## 2.0.0b9 (2026-07-28) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index e6efee6306ef..9c6b0b8ecb6c 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -59,6 +59,36 @@ def _read_task_manager_shutdown_grace() -> float: return 25.0 +def _resilient_tasks_opted_in() -> bool: + """Return True iff the app has declared at least one durable task. + + The opt-in signal is a non-empty ``_REGISTERED_DESCRIPTORS`` registry, + which every durable decorator populates: ``@task`` and + ``@multi_turn_task`` both construct a :class:`~azure.ai.agentserver.core. + tasks.Task` whose ``__init__`` appends to the registry (``MultiTurnTask`` + wraps an inner ``Task``, so it registers too). Declaring a durable task + is therefore the opt-in for standing up the ``TaskManager`` and its + (potentially network-backed) recovery scan. + + Returns False — so ``AgentServerHost`` skips ``TaskManager`` auto-init + entirely — when the resilient tasks module is unavailable or no durable + task was declared. This keeps plain servers (e.g. invocations-only + hosts that never use ``@task``) from paying the hosted task-store + startup cost: a blocking ``list()`` round-trip plus credential-token + acquisition that gates server readiness while having nothing to recover. + + :return: Whether resilient task auto-initialization should run. + :rtype: bool + """ + try: + from .tasks._decorator import ( # pylint: disable=import-outside-toplevel + _REGISTERED_DESCRIPTORS, + ) + except ImportError: + return False + return bool(_REGISTERED_DESCRIPTORS) + + def _mask_uri(uri: str) -> str: """Return only the scheme and host of a URI, hiding path/query/credentials. @@ -267,25 +297,37 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF ) # --- Resilient task manager auto-initialization --- + # + # OPT-IN GATE: only stand up the TaskManager when the app has + # actually declared a durable task via ``@task`` / + # ``@multi_turn_task`` (both funnel through the shared + # ``_REGISTERED_DESCRIPTORS`` registry). Apps that never opt in + # — e.g. plain invocations servers — must NOT pay the task-store + # cost: without this gate the startup path issues a blocking + # hosted task-store ``list()`` (plus DefaultAzureCredential token + # acquisition) BEFORE the lifespan ``yield``, gating server + # readiness on a network round-trip that has nothing to recover + # (no registered descriptors == no recovery-routing targets). task_manager = None - try: - from .tasks._manager import ( # pylint: disable=import-outside-toplevel - TaskManager, - set_task_manager, - ) + if _resilient_tasks_opted_in(): + try: + from .tasks._manager import ( # pylint: disable=import-outside-toplevel + TaskManager, + set_task_manager, + ) - task_manager = TaskManager( - config=cfg, - shutdown_event=asyncio.Event(), - shutdown_grace_seconds=_read_task_manager_shutdown_grace(), - ) - set_task_manager(task_manager) - await task_manager.startup() - logger.info("TaskManager initialized automatically") - except ImportError: - pass # resilient module not available - except Exception: # pylint: disable=broad-exception-caught - logger.warning("Failed to initialize TaskManager", exc_info=True) + task_manager = TaskManager( + config=cfg, + shutdown_event=asyncio.Event(), + shutdown_grace_seconds=_read_task_manager_shutdown_grace(), + ) + set_task_manager(task_manager) + await task_manager.startup() + logger.info("TaskManager initialized automatically") + except ImportError: + pass # resilient module not available + except Exception: # pylint: disable=broad-exception-caught + logger.warning("Failed to initialize TaskManager", exc_info=True) yield diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py new file mode 100644 index 000000000000..829413493cd7 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -0,0 +1,193 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for opt-in gating of resilient ``TaskManager`` auto-initialization. + +``AgentServerHost`` must only stand up the ``TaskManager`` (and its +potentially network-backed startup recovery scan) when the application has +actually declared a durable task via ``@task`` / ``@multi_turn_task``. +Plain servers that never opt in — e.g. invocations-only hosts — must NOT +pay the hosted task-store startup cost (a blocking ``list()`` round-trip +plus credential-token acquisition that would gate server readiness while +having nothing to recover). + +The opt-in signal is a non-empty ``_REGISTERED_DESCRIPTORS`` registry. +Both durable decorators funnel through the single ``Task.__init__`` +registration site (``MultiTurnTask`` wraps an inner ``Task``), so declaring +either kind opts the app in. +""" +import logging + +import pytest + +from azure.ai.agentserver.core._base import _resilient_tasks_opted_in +from azure.ai.agentserver.core.tasks import ( + TaskContext, + TaskManagerNotInitialized, + multi_turn_task, + task, +) +from azure.ai.agentserver.core.tasks import _decorator as _decorator_mod +from azure.ai.agentserver.core.tasks._manager import ( + get_task_manager, + set_task_manager, +) + + +@pytest.fixture +def _isolate_registry(): + """Snapshot/clear/restore the global durable-task registry + manager. + + ``_REGISTERED_DESCRIPTORS`` is process-global and populated at + decoration time by other test modules; snapshot and clear it so each + test controls the opt-in state, then restore it (and reset the manager + singleton) afterwards. + """ + saved = list(_decorator_mod._REGISTERED_DESCRIPTORS) + _decorator_mod._REGISTERED_DESCRIPTORS.clear() + set_task_manager(None) + try: + yield _decorator_mod._REGISTERED_DESCRIPTORS + finally: + _decorator_mod._REGISTERED_DESCRIPTORS.clear() + _decorator_mod._REGISTERED_DESCRIPTORS.extend(saved) + set_task_manager(None) + + +class _FakeTaskManager: + """Records construction + lifecycle without touching network/disk.""" + + instances: "list[_FakeTaskManager]" = [] + + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + self.startup_called = False + self.shutdown_called = False + _FakeTaskManager.instances.append(self) + + async def startup(self) -> None: + self.startup_called = True + + async def shutdown(self) -> None: + self.shutdown_called = True + + +@pytest.fixture +def _fake_task_manager(monkeypatch: pytest.MonkeyPatch): + """Patch the ``TaskManager`` the lifespan imports so no real provider, + token acquisition, or task-store round-trip happens during the test.""" + _FakeTaskManager.instances.clear() + monkeypatch.setattr( + "azure.ai.agentserver.core.tasks._manager.TaskManager", + _FakeTaskManager, + ) + return _FakeTaskManager + + +# ------------------------------------------------------------------ # +# _resilient_tasks_opted_in(): the opt-in signal +# ------------------------------------------------------------------ # + + +class TestOptInSignal: + """Unit tests for ``_resilient_tasks_opted_in()``.""" + + def test_empty_registry_is_not_opted_in(self, _isolate_registry) -> None: + assert _resilient_tasks_opted_in() is False + + def test_task_decorator_opts_in(self, _isolate_registry) -> None: + @task(name="optin_probe_one_shot") + async def _probe(ctx: "TaskContext[dict]") -> None: + return None + + assert _resilient_tasks_opted_in() is True + + def test_multi_turn_task_decorator_opts_in(self, _isolate_registry) -> None: + @multi_turn_task(name="optin_probe_multi_turn") + async def _probe(ctx: "TaskContext[dict]") -> None: + return None + + assert _resilient_tasks_opted_in() is True + + +# ------------------------------------------------------------------ # +# Lifespan behaviour: gate is honoured at startup +# ------------------------------------------------------------------ # + + +class TestLifespanOptInGate: + """Verify ``AgentServerHost`` lifespan honours the opt-in gate.""" + + @pytest.mark.asyncio + async def test_not_opted_in_skips_task_manager( + self, + _isolate_registry, + _fake_task_manager, + caplog: pytest.LogCaptureFixture, + ) -> None: + """No durable task declared → TaskManager is never constructed.""" + from azure.ai.agentserver.core import AgentServerHost + + app = AgentServerHost() + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + async with app.router.lifespan_context(app): + # During the active lifespan, no manager is installed. + with pytest.raises(TaskManagerNotInitialized): + get_task_manager() + + # TaskManager was never even constructed — no provider / token / + # task-store work happened. + assert _fake_task_manager.instances == [] + assert not any( + "TaskManager initialized automatically" in r.message for r in caplog.records + ) + + @pytest.mark.asyncio + async def test_opted_in_initializes_task_manager( + self, + _isolate_registry, + _fake_task_manager, + caplog: pytest.LogCaptureFixture, + ) -> None: + """A declared ``@task`` → TaskManager is constructed and started.""" + from azure.ai.agentserver.core import AgentServerHost + + @task(name="optin_lifespan_task") + async def _probe(ctx: "TaskContext[dict]") -> None: + return None + + app = AgentServerHost() + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + async with app.router.lifespan_context(app): + pass + + assert len(_fake_task_manager.instances) == 1 + mgr = _fake_task_manager.instances[0] + assert mgr.startup_called is True + assert mgr.shutdown_called is True + assert any( + "TaskManager initialized automatically" in r.message for r in caplog.records + ) + + @pytest.mark.asyncio + async def test_multi_turn_opt_in_initializes_task_manager( + self, + _isolate_registry, + _fake_task_manager, + ) -> None: + """A declared ``@multi_turn_task`` also opts the host in.""" + from azure.ai.agentserver.core import AgentServerHost + + @multi_turn_task(name="optin_lifespan_multi_turn") + async def _probe(ctx: "TaskContext[dict]") -> None: + return None + + app = AgentServerHost() + + async with app.router.lifespan_context(app): + pass + + assert len(_fake_task_manager.instances) == 1 + assert _fake_task_manager.instances[0].startup_called is True From bd26125944b72aa19d318d359291436177c508bf Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Thu, 30 Jul 2026 20:57:33 +0530 Subject: [PATCH 2/7] Preserve late-registration via lazy TaskManager bootstrap Address review: gating eager auto-init on a non-empty descriptor registry broke the supported late-registration path. Task.__init__ pushes decorators declared after startup into an already-live manager, but when the registry was empty at startup no manager existed, so a task declared during the active lifespan failed on first run/start with TaskManagerNotInitialized. Fix with lazy initialization instead of a hard skip: - The host always installs a lazy factory (set_task_manager_factory) that builds a TaskManager from config and loads registered descriptors. - Eager init + recovery scan still run when a durable task is declared by startup (unchanged behavior); otherwise nothing is constructed and the expensive startup path is deferred. - get_task_manager() lazily bootstraps via the factory on first task use, so a task declared mid-lifespan works instead of raising. - set_task_manager(None) also clears the factory; shutdown tears down whichever manager is live (eager or lazily bootstrapped) via a peek helper. - Extract TaskManager.load_registered_descriptors() reused by startup() and the factory. Outside a host lifespan (no factory) get_task_manager() still raises, so existing behavior is preserved. Adds coverage for declaring the first task inside the lifespan. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5222ab61-7d5e-4642-b75f-ef365e375f4f --- .../azure-ai-agentserver-core/CHANGELOG.md | 2 +- .../azure/ai/agentserver/core/_base.py | 87 +++++++++++++------ .../ai/agentserver/core/tasks/_manager.py | 84 +++++++++++++++--- .../tests/tasks/test_task_manager_optin.py | 61 +++++++++++-- 4 files changed, 189 insertions(+), 45 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 591d249b7d79..65b14ace90fc 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -8,7 +8,7 @@ ### Bugs Fixed -- `AgentServerHost` now only initializes the resilient `TaskManager` at startup when the app has declared a durable task via `@task` / `@multi_turn_task`. Plain servers (e.g. invocations-only hosts) that never use the task primitive no longer pay the task-store startup cost — a blocking hosted task-store `list()` round-trip plus credential-token acquisition that could gate server readiness (and add tens of seconds of connect latency) while having nothing to recover. +- `AgentServerHost` no longer eagerly initializes the resilient `TaskManager` at startup when the app has not declared a durable task via `@task` / `@multi_turn_task`. Plain servers (e.g. invocations-only hosts) that never use the task primitive no longer pay the task-store startup cost — a blocking hosted task-store `list()` round-trip plus credential-token acquisition that could gate server readiness (and add tens of seconds of connect latency) while having nothing to recover. A durable task declared *after* startup still works: its first `run`/`start` lazily bootstraps the `TaskManager` on demand instead of failing with `TaskManagerNotInitialized`. - Cleaned up `AgentServerHost` public signatures so inherited middleware typing does not expose Starlette private type aliases. ## 2.0.0b9 (2026-07-28) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index 9c6b0b8ecb6c..ef90bec04884 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -298,24 +298,44 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF # --- Resilient task manager auto-initialization --- # - # OPT-IN GATE: only stand up the TaskManager when the app has - # actually declared a durable task via ``@task`` / - # ``@multi_turn_task`` (both funnel through the shared - # ``_REGISTERED_DESCRIPTORS`` registry). Apps that never opt in - # — e.g. plain invocations servers — must NOT pay the task-store - # cost: without this gate the startup path issues a blocking - # hosted task-store ``list()`` (plus DefaultAzureCredential token - # acquisition) BEFORE the lifespan ``yield``, gating server - # readiness on a network round-trip that has nothing to recover - # (no registered descriptors == no recovery-routing targets). + # OPT-IN, with a lazy fallback for late registration. + # + # The TaskManager's startup recovery scan issues a blocking hosted + # task-store ``list()`` (plus DefaultAzureCredential token + # acquisition) that would otherwise gate server readiness on a + # network round-trip. Apps that never declare a durable task + # (``@task`` / ``@multi_turn_task`` — both funnel through the + # shared ``_REGISTERED_DESCRIPTORS`` registry) must not pay it. + # + # So: install a lazy factory unconditionally, then eagerly stand + # up the manager ONLY when a durable task was declared by startup. + # - Opted in -> eager init + recovery scan (unchanged behavior). + # - Not opted -> nothing is constructed now; if a task is later + # declared during the active lifespan (the late-registration + # path in ``Task.__init__``), its first ``run``/``start`` + # triggers ``get_task_manager()``, which builds the manager via + # this factory instead of raising ``TaskManagerNotInitialized``. task_manager = None - if _resilient_tasks_opted_in(): - try: - from .tasks._manager import ( # pylint: disable=import-outside-toplevel - TaskManager, - set_task_manager, + try: + from .tasks._manager import ( # pylint: disable=import-outside-toplevel + TaskManager, + set_task_manager, + set_task_manager_factory, + ) + + def _bootstrap_task_manager() -> "TaskManager": + manager = TaskManager( + config=cfg, + shutdown_event=asyncio.Event(), + shutdown_grace_seconds=_read_task_manager_shutdown_grace(), ) + manager.load_registered_descriptors() + logger.info("TaskManager lazily initialized for a task declared after startup") + return manager + set_task_manager_factory(_bootstrap_task_manager) + + if _resilient_tasks_opted_in(): task_manager = TaskManager( config=cfg, shutdown_event=asyncio.Event(), @@ -324,10 +344,15 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF set_task_manager(task_manager) await task_manager.startup() logger.info("TaskManager initialized automatically") - except ImportError: - pass # resilient module not available - except Exception: # pylint: disable=broad-exception-caught - logger.warning("Failed to initialize TaskManager", exc_info=True) + else: + logger.info( + "TaskManager auto-init deferred (no durable task declared); " + "will lazily initialize on first task use" + ) + except ImportError: + pass # resilient module not available + except Exception: # pylint: disable=broad-exception-caught + logger.warning("Failed to initialize TaskManager", exc_info=True) yield @@ -362,18 +387,28 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF # Shutdown task manager AFTER on_shutdown so resilient handlers # have had time to checkpoint via the responses layer's - # ``handle_shutdown``. - if task_manager is not None: - try: - await task_manager.shutdown() - from .tasks._manager import ( # pylint: disable=import-outside-toplevel - set_task_manager as _clear_manager, - ) + # ``handle_shutdown``. Tear down whichever manager is live — + # the eagerly-initialized one OR one lazily bootstrapped by a + # task declared during the active lifespan — via the peek helper + # (so a lazily-created manager not captured in ``task_manager`` + # is still shut down and cleared). + from .tasks._manager import ( # pylint: disable=import-outside-toplevel + _peek_task_manager, + set_task_manager as _clear_manager, + ) + live_task_manager = task_manager or _peek_task_manager() + if live_task_manager is not None: + try: + await live_task_manager.shutdown() _clear_manager(None) logger.info("TaskManager shut down") except Exception: # pylint: disable=broad-exception-caught logger.warning("Error shutting down TaskManager", exc_info=True) + else: + # No manager was ever created; clear the lazy factory so the + # global state doesn't leak past this lifespan. + _clear_manager(None) # Merge routes: subclass routes (if any) + health endpoint all_routes: list[Any] = list(routes or []) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py index f73c226d20f9..9a6fc01b656d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py @@ -78,6 +78,15 @@ # Module-level manager singleton _manager: TaskManager | None = None +# Optional zero-arg factory that lazily builds (and registers) a TaskManager +# on first use. Installed by ``AgentServerHost`` during lifespan startup so a +# durable task declared *after* startup (the late-registration path in +# ``Task.__init__``) can still bootstrap a manager on its first ``run``/ +# ``start`` — even when no task was declared at startup and eager auto-init +# was therefore skipped. ``None`` outside an active host lifespan (e.g. unit +# tests), in which case ``get_task_manager`` raises as before. +_manager_factory: "Callable[[], TaskManager] | None" = None + def _is_evicted(exc: BaseException) -> bool: """Return True if ``exc`` is the eviction-classified rejection. @@ -265,12 +274,23 @@ def _lease_is_dead( def get_task_manager() -> TaskManager: """Return the active TaskManager singleton. + If no manager has been installed but a lazy factory has been registered + (by ``AgentServerHost`` during lifespan startup), the manager is + constructed on demand and cached. This preserves the late-registration + path: a durable task declared *after* startup can still bootstrap the + manager on its first ``run``/``start`` even when eager auto-init was + skipped because no task was declared at startup. + :raises ~azure.ai.agentserver.core.tasks.TaskManagerNotInitialized: If no - manager has been initialized. + manager has been initialized and no lazy factory is available. :return: The active manager. :rtype: TaskManager """ + global _manager # pylint: disable=global-statement if _manager is None: + if _manager_factory is not None: + _manager = _manager_factory() + return _manager from ._exceptions import ( # pylint: disable=import-outside-toplevel TaskManagerNotInitialized, ) @@ -282,16 +302,49 @@ def get_task_manager() -> TaskManager: return _manager +def _peek_task_manager() -> "TaskManager | None": + """Return the installed manager without lazily creating one (or raising). + + Used by the host shutdown path to tear down whichever manager is live — + whether it was eagerly initialized at startup or lazily bootstrapped by a + post-startup task registration. + + :return: The current manager, or ``None`` if none is installed. + :rtype: TaskManager | None + """ + return _manager + + +def set_task_manager_factory(factory: "Callable[[], TaskManager] | None") -> None: + """Install (or clear) the lazy manager factory. + + Called by ``AgentServerHost`` during lifespan startup with a factory that + builds a manager and loads the registered task descriptors into it. The + factory is invoked at most once, by :func:`get_task_manager`, when a task + is first used without an eagerly-initialized manager. + + :param factory: A zero-arg callable returning a ready ``TaskManager``, or + ``None`` to clear. + :type factory: Callable[[], TaskManager] | None + """ + global _manager_factory # pylint: disable=global-statement + _manager_factory = factory + + def set_task_manager(manager: TaskManager | None) -> None: """Set the module-level TaskManager singleton. - Called by ``AgentServerHost`` during startup/shutdown. + Called by ``AgentServerHost`` during startup/shutdown. Clearing the + manager (``manager=None``) also clears any installed lazy factory, so a + full reset leaves neither a live manager nor a pending lazy bootstrap. :param manager: The manager to set, or ``None`` to clear. :type manager: TaskManager | None """ - global _manager # pylint: disable=global-statement + global _manager, _manager_factory # pylint: disable=global-statement _manager = manager + if manager is None: + _manager_factory = None class _ActiveTask: # pylint: disable=too-many-instance-attributes @@ -704,6 +757,23 @@ async def _cancel_queued_steering_input( # pylint: disable=unused-argument if not future.done(): future.set_exception(TaskCancelled()) + def load_registered_descriptors(self) -> None: + """Load module-level ``@task`` / ``@multi_turn_task`` descriptors. + + Populates ``_resume_callbacks`` / ``_resume_opts`` from the global + ``_REGISTERED_DESCRIPTORS`` registry so recovery routing and + multi-turn opts resolution can find the right handler. Called by + :meth:`startup` and by the lazy-bootstrap factory so a manager built + after startup still picks up every declared task. + """ + from ._decorator import ( # pylint: disable=import-outside-toplevel + _REGISTERED_DESCRIPTORS, + ) + + for fn_name, fn, opts in _REGISTERED_DESCRIPTORS: + self._resume_callbacks[fn_name] = fn + self._resume_opts[fn_name] = opts + async def startup(self) -> None: """Initialize the manager and recover stale tasks. @@ -716,13 +786,7 @@ async def startup(self) -> None: self._config.is_hosted, ) # Pick up descriptors registered at import time (for recovery) - from ._decorator import ( # pylint: disable=import-outside-toplevel - _REGISTERED_DESCRIPTORS, - ) - - for fn_name, fn, opts in _REGISTERED_DESCRIPTORS: - self._resume_callbacks[fn_name] = fn - self._resume_opts[fn_name] = opts + self.load_registered_descriptors() await self._recover_stale_tasks() diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py index 829413493cd7..06635f9f8ddf 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -63,8 +63,12 @@ def __init__(self, **kwargs) -> None: self.kwargs = kwargs self.startup_called = False self.shutdown_called = False + self.descriptors_loaded = False _FakeTaskManager.instances.append(self) + def load_registered_descriptors(self) -> None: + self.descriptors_loaded = True + async def startup(self) -> None: self.startup_called = True @@ -109,6 +113,14 @@ async def _probe(ctx: "TaskContext[dict]") -> None: assert _resilient_tasks_opted_in() is True + def test_get_task_manager_still_raises_without_a_factory( + self, _isolate_registry + ) -> None: + """Outside an active host lifespan (no factory installed), the raise + behavior is preserved — lazy bootstrap only applies within a host.""" + with pytest.raises(TaskManagerNotInitialized): + get_task_manager() + # ------------------------------------------------------------------ # # Lifespan behaviour: gate is honoured at startup @@ -119,29 +131,62 @@ class TestLifespanOptInGate: """Verify ``AgentServerHost`` lifespan honours the opt-in gate.""" @pytest.mark.asyncio - async def test_not_opted_in_skips_task_manager( + async def test_not_opted_in_defers_eager_init( self, _isolate_registry, _fake_task_manager, caplog: pytest.LogCaptureFixture, ) -> None: - """No durable task declared → TaskManager is never constructed.""" + """No durable task declared → no manager is eagerly constructed and + no recovery scan runs; the expensive startup path is deferred.""" from azure.ai.agentserver.core import AgentServerHost app = AgentServerHost() with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): async with app.router.lifespan_context(app): - # During the active lifespan, no manager is installed. - with pytest.raises(TaskManagerNotInitialized): - get_task_manager() + # Nothing eagerly built during startup — no provider / token / + # task-store round-trip happened. + assert _fake_task_manager.instances == [] - # TaskManager was never even constructed — no provider / token / - # task-store work happened. - assert _fake_task_manager.instances == [] assert not any( "TaskManager initialized automatically" in r.message for r in caplog.records ) + assert any( + "TaskManager auto-init deferred" in r.message for r in caplog.records + ) + + @pytest.mark.asyncio + async def test_first_task_declared_in_lifespan_lazily_bootstraps( + self, + _isolate_registry, + _fake_task_manager, + ) -> None: + """Late-registration path: a task declared *during* the active + lifespan (when eager init was skipped) must not fail on first use — + ``get_task_manager()`` lazily bootstraps a manager instead of raising + ``TaskManagerNotInitialized``.""" + from azure.ai.agentserver.core import AgentServerHost + + app = AgentServerHost() + + async with app.router.lifespan_context(app): + # Nothing built yet. + assert _fake_task_manager.instances == [] + + # Declare the FIRST durable task now, mid-lifespan. + @task(name="optin_late_declared_task") + async def _late(ctx: "TaskContext[dict]") -> None: + return None + + # First use resolves the manager — must not raise, and must + # lazily construct + load descriptors. + mgr = get_task_manager() + assert mgr is _fake_task_manager.instances[0] + assert mgr.descriptors_loaded is True + + # The lazily-bootstrapped manager is torn down on shutdown. + assert _fake_task_manager.instances[0].shutdown_called is True @pytest.mark.asyncio async def test_opted_in_initializes_task_manager( From 9282136cd247cbad3620bb40b2b7d1ee909448b8 Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Thu, 30 Jul 2026 21:20:48 +0530 Subject: [PATCH 3/7] Guard shutdown import of tasks module; harden lazy-bootstrap test Address review round 2 (suppressed low-confidence findings): - Real bug: the lifespan shutdown path imported .tasks._manager without an ImportError guard, while startup tolerates it. In a task-less/optional install, lifespan entry would succeed but teardown would raise ImportError. Guard the shutdown import the same way (try/except ImportError) so optional installations stop cleanly. Add a test that blocks the module import and asserts both lifespan entry and exit succeed. - Make the late-registration test honest about the intentional limitation: assert the lazily-bootstrapped manager's startup() is NOT called (no initial recovery scan / periodic loop), documenting current behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5222ab61-7d5e-4642-b75f-ef365e375f4f --- .../azure/ai/agentserver/core/_base.py | 39 +++++++++++-------- .../tests/tasks/test_task_manager_optin.py | 37 +++++++++++++++++- 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index ef90bec04884..fe37c8722852 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -391,24 +391,29 @@ def _bootstrap_task_manager() -> "TaskManager": # the eagerly-initialized one OR one lazily bootstrapped by a # task declared during the active lifespan — via the peek helper # (so a lazily-created manager not captured in ``task_manager`` - # is still shut down and cleared). - from .tasks._manager import ( # pylint: disable=import-outside-toplevel - _peek_task_manager, - set_task_manager as _clear_manager, - ) - - live_task_manager = task_manager or _peek_task_manager() - if live_task_manager is not None: - try: - await live_task_manager.shutdown() - _clear_manager(None) - logger.info("TaskManager shut down") - except Exception: # pylint: disable=broad-exception-caught - logger.warning("Error shutting down TaskManager", exc_info=True) + # is still shut down and cleared). Guard the import with the same + # ``ImportError`` tolerance as startup: task-less / optional + # installations must still tear down cleanly. + try: + from .tasks._manager import ( # pylint: disable=import-outside-toplevel + _peek_task_manager, + set_task_manager as _clear_manager, + ) + except ImportError: + pass # resilient module not available — nothing to tear down else: - # No manager was ever created; clear the lazy factory so the - # global state doesn't leak past this lifespan. - _clear_manager(None) + live_task_manager = task_manager or _peek_task_manager() + if live_task_manager is not None: + try: + await live_task_manager.shutdown() + _clear_manager(None) + logger.info("TaskManager shut down") + except Exception: # pylint: disable=broad-exception-caught + logger.warning("Error shutting down TaskManager", exc_info=True) + else: + # No manager was ever created; clear the lazy factory so + # the global state doesn't leak past this lifespan. + _clear_manager(None) # Merge routes: subclass routes (if any) + health endpoint all_routes: list[Any] = list(routes or []) diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py index 06635f9f8ddf..b4d7742527e0 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -180,10 +180,17 @@ async def _late(ctx: "TaskContext[dict]") -> None: return None # First use resolves the manager — must not raise, and must - # lazily construct + load descriptors. + # lazily construct + load descriptors. This is exactly the first + # thing ``Task.run``/``Task.start`` do, so it protects the + # advertised "first run/start no longer raises" behavior. mgr = get_task_manager() assert mgr is _fake_task_manager.instances[0] assert mgr.descriptors_loaded is True + # Intentional limitation: the sync lazy path does NOT run the + # async ``startup()`` (initial recovery scan + periodic loop). + # A task declared after startup already missed the startup + # recovery window; its own run/start works via the provider. + assert mgr.startup_called is False # The lazily-bootstrapped manager is torn down on shutdown. assert _fake_task_manager.instances[0].shutdown_called is True @@ -236,3 +243,31 @@ async def _probe(ctx: "TaskContext[dict]") -> None: assert len(_fake_task_manager.instances) == 1 assert _fake_task_manager.instances[0].startup_called is True + + @pytest.mark.asyncio + async def test_lifespan_tears_down_cleanly_when_tasks_module_unavailable( + self, + _isolate_registry, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Task-less / optional install: both startup AND shutdown tolerate + the resilient-task module being unimportable (no crash on teardown).""" + import builtins + + real_import = builtins.__import__ + + def _blocked_import(name, *args, **kwargs): + if name.endswith("tasks._manager") or name == "azure.ai.agentserver.core.tasks._manager": + raise ImportError("simulated: resilient tasks module unavailable") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _blocked_import) + + from azure.ai.agentserver.core import AgentServerHost + + app = AgentServerHost() + + # Entering AND exiting the lifespan must both succeed without raising + # ImportError — the shutdown import is guarded like the startup one. + async with app.router.lifespan_context(app): + pass From a7b1a0238c08d83706adbfd4c099f9cc2bd667d5 Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Fri, 31 Jul 2026 00:21:32 +0530 Subject: [PATCH 4/7] Gate TaskManager on explicit switch AND declared tasks; drop lazy factory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the registry-only opt-in with a double gate per reviewer feedback (Shivakishore14: _REGISTERED_DESCRIPTORS alone is a process-global heuristic prone to false positives). AgentServerHost now stands up the resilient TaskManager (and its network-backed startup recovery scan) only when BOTH: 1. set_resilient_tasks_enabled(True) was called (new process-global switch, default False), AND 2. the _REGISTERED_DESCRIPTORS list is non-empty (a durable task declared). Both signals are read directly at lifespan startup — no lazy factory, no opted-in abstraction. Adds set_resilient_tasks_enabled / resilient_tasks_enabled to the tasks public API. Reverts the _manager.py factory machinery from earlier commits (back to origin/main); the double gate makes it unnecessary. - New tasks/_enablement.py: process-global switch (default disabled). - _base.py: gate = _resilient_tasks_enabled() and _has_registered_tasks(). - Tests rewritten for the AND truth table (4 lifespan cases + signal units). - Updated public-API-surface expected set, tasks-guide.md, api.md, CHANGELOG. Note: the responses layer (uses tasks internally) must call set_resilient_tasks_enabled(True) in its host — tracked for PR #47275. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5222ab61-7d5e-4642-b75f-ef365e375f4f --- .../azure-ai-agentserver-core/CHANGELOG.md | 3 +- .../azure-ai-agentserver-core/api.md | 6 + .../api.metadata.yml | 4 +- .../azure/ai/agentserver/core/_base.py | 160 +++++------ .../ai/agentserver/core/tasks/__init__.py | 4 + .../ai/agentserver/core/tasks/_enablement.py | 47 ++++ .../ai/agentserver/core/tasks/_manager.py | 84 +----- .../docs/tasks-guide.md | 22 ++ .../tests/tasks/test_public_api_surface.py | 3 + .../tests/tasks/test_task_manager_optin.py | 254 +++++++----------- 10 files changed, 262 insertions(+), 325 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 65b14ace90fc..71407e8058cb 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -5,10 +5,11 @@ ### Features Added - Added public `MiddlewareFactory` and `StreamContent` typing aliases for host middleware and streaming helpers. +- Added `set_resilient_tasks_enabled` / `resilient_tasks_enabled` to `azure.ai.agentserver.core.tasks` — an explicit process-global switch (default **disabled**) that gates the resilient `TaskManager`. ### Bugs Fixed -- `AgentServerHost` no longer eagerly initializes the resilient `TaskManager` at startup when the app has not declared a durable task via `@task` / `@multi_turn_task`. Plain servers (e.g. invocations-only hosts) that never use the task primitive no longer pay the task-store startup cost — a blocking hosted task-store `list()` round-trip plus credential-token acquisition that could gate server readiness (and add tens of seconds of connect latency) while having nothing to recover. A durable task declared *after* startup still works: its first `run`/`start` lazily bootstraps the `TaskManager` on demand instead of failing with `TaskManagerNotInitialized`. +- `AgentServerHost` now initializes the resilient `TaskManager` (and its network-backed startup recovery scan) only when **both** the subsystem was explicitly enabled via `set_resilient_tasks_enabled(True)` **and** at least one durable task (`@task` / `@multi_turn_task`) has been declared. Plain servers (e.g. invocations-only hosts) that do not opt in no longer pay the task-store startup cost — a blocking hosted task-store `list()` round-trip plus credential-token acquisition that could gate server readiness and add tens of seconds of connect latency while having nothing to recover. - Cleaned up `AgentServerHost` public signatures so inherited middleware typing does not expose Starlette private type aliases. ## 2.0.0b9 (2026-07-28) diff --git a/sdk/agentserver/azure-ai-agentserver-core/api.md b/sdk/agentserver/azure-ai-agentserver-core/api.md index 7fa3dee2885a..4e174496469d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/api.md +++ b/sdk/agentserver/azure-ai-agentserver-core/api.md @@ -591,6 +591,12 @@ namespace azure.ai.agentserver.core.tasks ) -> Callable[[Callable[[TaskContext[Input]], Awaitable[Output]]], MultiTurnTask[Input, Output]]: ... + def azure.ai.agentserver.core.tasks.resilient_tasks_enabled() -> bool: ... + + + def azure.ai.agentserver.core.tasks.set_resilient_tasks_enabled(value: bool = True) -> None: ... + + @overload def azure.ai.agentserver.core.tasks.task( fn: Callable[[TaskContext[Input]], Awaitable[Output]], diff --git a/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml b/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml index a2bb5463859b..1ffcb576c382 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml +++ b/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml @@ -1,3 +1,3 @@ -apiMdSha256: 803cc6578b722fdf575384cd50fee9429f690be7e4036cd805ff4fa9090892f8 +apiMdSha256: fd3b7bd1ae87d9b18719027594934c55d968973f159fbc8a5b416d227c67282c parserVersion: 0.3.30 -pythonVersion: 3.14.3 +pythonVersion: 3.12.10 diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index fe37c8722852..ea785be98855 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -59,25 +59,20 @@ def _read_task_manager_shutdown_grace() -> float: return 25.0 -def _resilient_tasks_opted_in() -> bool: - """Return True iff the app has declared at least one durable task. - - The opt-in signal is a non-empty ``_REGISTERED_DESCRIPTORS`` registry, - which every durable decorator populates: ``@task`` and - ``@multi_turn_task`` both construct a :class:`~azure.ai.agentserver.core. - tasks.Task` whose ``__init__`` appends to the registry (``MultiTurnTask`` - wraps an inner ``Task``, so it registers too). Declaring a durable task - is therefore the opt-in for standing up the ``TaskManager`` and its - (potentially network-backed) recovery scan. - - Returns False — so ``AgentServerHost`` skips ``TaskManager`` auto-init - entirely — when the resilient tasks module is unavailable or no durable - task was declared. This keeps plain servers (e.g. invocations-only - hosts that never use ``@task``) from paying the hosted task-store - startup cost: a blocking ``list()`` round-trip plus credential-token - acquisition that gates server readiness while having nothing to recover. - - :return: Whether resilient task auto-initialization should run. +def _has_registered_tasks() -> bool: + """Return True iff the ``_REGISTERED_DESCRIPTORS`` list is non-empty. + + That list is populated at import time by the durable-task decorators: + ``@task`` and ``@multi_turn_task`` both construct a + :class:`~azure.ai.agentserver.core.tasks.Task` whose ``__init__`` appends + to it (``MultiTurnTask`` wraps an inner ``Task``, so it registers too). It + is the SAME list ``TaskManager.startup()`` already reads to bind recovery + callbacks, so checking it here introduces no new timing assumption. + + Returns False when the resilient tasks module is unavailable or no durable + task has been declared. + + :return: Whether at least one durable task is registered. :rtype: bool """ try: @@ -89,6 +84,26 @@ def _resilient_tasks_opted_in() -> bool: return bool(_REGISTERED_DESCRIPTORS) +def _resilient_tasks_enabled() -> bool: + """Return True iff the resilient task subsystem was explicitly enabled. + + Reads the process-global switch toggled by + :func:`~azure.ai.agentserver.core.tasks.set_resilient_tasks_enabled` + (defaults to ``False``). Returns False when the resilient tasks module is + unavailable. + + :return: Whether resilient tasks were explicitly enabled. + :rtype: bool + """ + try: + from .tasks._enablement import ( # pylint: disable=import-outside-toplevel + resilient_tasks_enabled, + ) + except ImportError: + return False + return resilient_tasks_enabled() + + def _mask_uri(uri: str) -> str: """Return only the scheme and host of a URI, hiding path/query/credentials. @@ -298,44 +313,27 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF # --- Resilient task manager auto-initialization --- # - # OPT-IN, with a lazy fallback for late registration. - # - # The TaskManager's startup recovery scan issues a blocking hosted - # task-store ``list()`` (plus DefaultAzureCredential token - # acquisition) that would otherwise gate server readiness on a - # network round-trip. Apps that never declare a durable task - # (``@task`` / ``@multi_turn_task`` — both funnel through the - # shared ``_REGISTERED_DESCRIPTORS`` registry) must not pay it. + # DOUBLE GATE (AND): the TaskManager — and its network-backed + # startup recovery scan (a blocking hosted task-store ``list()`` + # plus DefaultAzureCredential token acquisition) — is stood up ONLY + # when BOTH hold: + # (1) resilient tasks were explicitly enabled via + # ``set_resilient_tasks_enabled(True)`` (default False), AND + # (2) the ``_REGISTERED_DESCRIPTORS`` list is non-empty, i.e. at + # least one durable task (``@task`` / ``@multi_turn_task``) + # has been declared. # - # So: install a lazy factory unconditionally, then eagerly stand - # up the manager ONLY when a durable task was declared by startup. - # - Opted in -> eager init + recovery scan (unchanged behavior). - # - Not opted -> nothing is constructed now; if a task is later - # declared during the active lifespan (the late-registration - # path in ``Task.__init__``), its first ``run``/``start`` - # triggers ``get_task_manager()``, which builds the manager via - # this factory instead of raising ``TaskManagerNotInitialized``. + # Both are read directly here. If either is false, nothing is + # constructed and no task-store call is made — plain servers (e.g. + # invocations-only hosts) pay nothing. task_manager = None - try: - from .tasks._manager import ( # pylint: disable=import-outside-toplevel - TaskManager, - set_task_manager, - set_task_manager_factory, - ) - - def _bootstrap_task_manager() -> "TaskManager": - manager = TaskManager( - config=cfg, - shutdown_event=asyncio.Event(), - shutdown_grace_seconds=_read_task_manager_shutdown_grace(), + if _resilient_tasks_enabled() and _has_registered_tasks(): + try: + from .tasks._manager import ( # pylint: disable=import-outside-toplevel + TaskManager, + set_task_manager, ) - manager.load_registered_descriptors() - logger.info("TaskManager lazily initialized for a task declared after startup") - return manager - set_task_manager_factory(_bootstrap_task_manager) - - if _resilient_tasks_opted_in(): task_manager = TaskManager( config=cfg, shutdown_event=asyncio.Event(), @@ -344,15 +342,16 @@ def _bootstrap_task_manager() -> "TaskManager": set_task_manager(task_manager) await task_manager.startup() logger.info("TaskManager initialized automatically") - else: - logger.info( - "TaskManager auto-init deferred (no durable task declared); " - "will lazily initialize on first task use" - ) - except ImportError: - pass # resilient module not available - except Exception: # pylint: disable=broad-exception-caught - logger.warning("Failed to initialize TaskManager", exc_info=True) + except ImportError: + pass # resilient module not available + except Exception: # pylint: disable=broad-exception-caught + logger.warning("Failed to initialize TaskManager", exc_info=True) + else: + logger.info( + "TaskManager not initialized (enabled=%s, tasks_declared=%s)", + _resilient_tasks_enabled(), + _has_registered_tasks(), + ) yield @@ -387,33 +386,18 @@ def _bootstrap_task_manager() -> "TaskManager": # Shutdown task manager AFTER on_shutdown so resilient handlers # have had time to checkpoint via the responses layer's - # ``handle_shutdown``. Tear down whichever manager is live — - # the eagerly-initialized one OR one lazily bootstrapped by a - # task declared during the active lifespan — via the peek helper - # (so a lazily-created manager not captured in ``task_manager`` - # is still shut down and cleared). Guard the import with the same - # ``ImportError`` tolerance as startup: task-less / optional - # installations must still tear down cleanly. - try: - from .tasks._manager import ( # pylint: disable=import-outside-toplevel - _peek_task_manager, - set_task_manager as _clear_manager, - ) - except ImportError: - pass # resilient module not available — nothing to tear down - else: - live_task_manager = task_manager or _peek_task_manager() - if live_task_manager is not None: - try: - await live_task_manager.shutdown() - _clear_manager(None) - logger.info("TaskManager shut down") - except Exception: # pylint: disable=broad-exception-caught - logger.warning("Error shutting down TaskManager", exc_info=True) - else: - # No manager was ever created; clear the lazy factory so - # the global state doesn't leak past this lifespan. + # ``handle_shutdown``. + if task_manager is not None: + try: + await task_manager.shutdown() + from .tasks._manager import ( # pylint: disable=import-outside-toplevel + set_task_manager as _clear_manager, + ) + _clear_manager(None) + logger.info("TaskManager shut down") + except Exception: # pylint: disable=broad-exception-caught + logger.warning("Error shutting down TaskManager", exc_info=True) # Merge routes: subclass routes (if any) + health endpoint all_routes: list[Any] = list(routes or []) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/__init__.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/__init__.py index f7956dd43b3f..38ad9359e2e7 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/__init__.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/__init__.py @@ -49,6 +49,7 @@ from ._context import EntryMode, TaskContext from ._decorator import MultiTurnTask, Task, multi_turn_task, task +from ._enablement import resilient_tasks_enabled, set_resilient_tasks_enabled from ._exceptions import ( InputTooLarge, LastInputIdPreconditionFailed, @@ -82,6 +83,9 @@ "multi_turn_task", "Task", "MultiTurnTask", + # Enablement switch + "set_resilient_tasks_enabled", + "resilient_tasks_enabled", # Context + metadata "TaskContext", "TaskMetadata", diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py new file mode 100644 index 000000000000..eadc355b9f4e --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py @@ -0,0 +1,47 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Explicit enablement switch for the resilient task subsystem. + +The resilient ``TaskManager`` performs network-backed startup work (a hosted +task-store recovery scan plus credential-token acquisition). To avoid making +any of those calls unless an app genuinely relies on durable tasks, +``AgentServerHost`` only stands up the ``TaskManager`` when BOTH conditions +hold: + +1. this switch has been turned on via :func:`set_resilient_tasks_enabled`, AND +2. at least one durable task has been declared (``@task`` / ``@multi_turn_task``). + +The switch is process-global and defaults to ``False`` (disabled). It is +intentionally decoupled from any ``AgentServerHost`` instance so it can be +flipped independently at import time, e.g.:: + + from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled + + set_resilient_tasks_enabled(True) +""" + +_RESILIENT_TASKS_ENABLED: bool = False + + +def set_resilient_tasks_enabled(value: bool = True) -> None: + """Enable or disable the resilient task subsystem process-wide. + + Must be called before ``AgentServerHost`` lifespan startup (typically at + import time) to take effect for eager initialization. Defaults the switch + to enabled when called with no argument. + + :param value: ``True`` to enable resilient tasks, ``False`` to disable. + :type value: bool + """ + global _RESILIENT_TASKS_ENABLED # pylint: disable=global-statement + _RESILIENT_TASKS_ENABLED = bool(value) + + +def resilient_tasks_enabled() -> bool: + """Return whether the resilient task subsystem has been explicitly enabled. + + :return: ``True`` if :func:`set_resilient_tasks_enabled` turned it on. + :rtype: bool + """ + return _RESILIENT_TASKS_ENABLED diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py index 9a6fc01b656d..f73c226d20f9 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py @@ -78,15 +78,6 @@ # Module-level manager singleton _manager: TaskManager | None = None -# Optional zero-arg factory that lazily builds (and registers) a TaskManager -# on first use. Installed by ``AgentServerHost`` during lifespan startup so a -# durable task declared *after* startup (the late-registration path in -# ``Task.__init__``) can still bootstrap a manager on its first ``run``/ -# ``start`` — even when no task was declared at startup and eager auto-init -# was therefore skipped. ``None`` outside an active host lifespan (e.g. unit -# tests), in which case ``get_task_manager`` raises as before. -_manager_factory: "Callable[[], TaskManager] | None" = None - def _is_evicted(exc: BaseException) -> bool: """Return True if ``exc`` is the eviction-classified rejection. @@ -274,23 +265,12 @@ def _lease_is_dead( def get_task_manager() -> TaskManager: """Return the active TaskManager singleton. - If no manager has been installed but a lazy factory has been registered - (by ``AgentServerHost`` during lifespan startup), the manager is - constructed on demand and cached. This preserves the late-registration - path: a durable task declared *after* startup can still bootstrap the - manager on its first ``run``/``start`` even when eager auto-init was - skipped because no task was declared at startup. - :raises ~azure.ai.agentserver.core.tasks.TaskManagerNotInitialized: If no - manager has been initialized and no lazy factory is available. + manager has been initialized. :return: The active manager. :rtype: TaskManager """ - global _manager # pylint: disable=global-statement if _manager is None: - if _manager_factory is not None: - _manager = _manager_factory() - return _manager from ._exceptions import ( # pylint: disable=import-outside-toplevel TaskManagerNotInitialized, ) @@ -302,49 +282,16 @@ def get_task_manager() -> TaskManager: return _manager -def _peek_task_manager() -> "TaskManager | None": - """Return the installed manager without lazily creating one (or raising). - - Used by the host shutdown path to tear down whichever manager is live — - whether it was eagerly initialized at startup or lazily bootstrapped by a - post-startup task registration. - - :return: The current manager, or ``None`` if none is installed. - :rtype: TaskManager | None - """ - return _manager - - -def set_task_manager_factory(factory: "Callable[[], TaskManager] | None") -> None: - """Install (or clear) the lazy manager factory. - - Called by ``AgentServerHost`` during lifespan startup with a factory that - builds a manager and loads the registered task descriptors into it. The - factory is invoked at most once, by :func:`get_task_manager`, when a task - is first used without an eagerly-initialized manager. - - :param factory: A zero-arg callable returning a ready ``TaskManager``, or - ``None`` to clear. - :type factory: Callable[[], TaskManager] | None - """ - global _manager_factory # pylint: disable=global-statement - _manager_factory = factory - - def set_task_manager(manager: TaskManager | None) -> None: """Set the module-level TaskManager singleton. - Called by ``AgentServerHost`` during startup/shutdown. Clearing the - manager (``manager=None``) also clears any installed lazy factory, so a - full reset leaves neither a live manager nor a pending lazy bootstrap. + Called by ``AgentServerHost`` during startup/shutdown. :param manager: The manager to set, or ``None`` to clear. :type manager: TaskManager | None """ - global _manager, _manager_factory # pylint: disable=global-statement + global _manager # pylint: disable=global-statement _manager = manager - if manager is None: - _manager_factory = None class _ActiveTask: # pylint: disable=too-many-instance-attributes @@ -757,23 +704,6 @@ async def _cancel_queued_steering_input( # pylint: disable=unused-argument if not future.done(): future.set_exception(TaskCancelled()) - def load_registered_descriptors(self) -> None: - """Load module-level ``@task`` / ``@multi_turn_task`` descriptors. - - Populates ``_resume_callbacks`` / ``_resume_opts`` from the global - ``_REGISTERED_DESCRIPTORS`` registry so recovery routing and - multi-turn opts resolution can find the right handler. Called by - :meth:`startup` and by the lazy-bootstrap factory so a manager built - after startup still picks up every declared task. - """ - from ._decorator import ( # pylint: disable=import-outside-toplevel - _REGISTERED_DESCRIPTORS, - ) - - for fn_name, fn, opts in _REGISTERED_DESCRIPTORS: - self._resume_callbacks[fn_name] = fn - self._resume_opts[fn_name] = opts - async def startup(self) -> None: """Initialize the manager and recover stale tasks. @@ -786,7 +716,13 @@ async def startup(self) -> None: self._config.is_hosted, ) # Pick up descriptors registered at import time (for recovery) - self.load_registered_descriptors() + from ._decorator import ( # pylint: disable=import-outside-toplevel + _REGISTERED_DESCRIPTORS, + ) + + for fn_name, fn, opts in _REGISTERED_DESCRIPTORS: + self._resume_callbacks[fn_name] = fn + self._resume_opts[fn_name] = opts await self._recover_stale_tasks() diff --git a/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md b/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md index d0d1b8e135dc..0a8f99d3574e 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md @@ -113,6 +113,28 @@ What this primitive deliberately does **not** do: ## 3. Hello world +### Enabling resilient tasks + +The resilient `TaskManager` — and its startup recovery scan against the hosted +task store — is **opt-in**. `AgentServerHost` only stands it up when **both** +of the following are true: + +1. the subsystem was explicitly enabled via `set_resilient_tasks_enabled(True)` + (it defaults to disabled), **and** +2. at least one durable task has been declared (`@task` / `@multi_turn_task`). + +Call `set_resilient_tasks_enabled(True)` once, before the host starts (e.g. at +module import). Use `resilient_tasks_enabled()` to read the current state. + +```python +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled + +set_resilient_tasks_enabled(True) # required for the TaskManager to run +``` + +Servers that never declare a task — or never enable the switch — skip the +TaskManager entirely and pay none of its startup cost. + ### One-shot ```python diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_public_api_surface.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_public_api_surface.py index 1437bcb113db..4ba14fa9cc4d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_public_api_surface.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_public_api_surface.py @@ -40,6 +40,9 @@ "multi_turn_task", "Task", "MultiTurnTask", + # Enablement switch + "set_resilient_tasks_enabled", + "resilient_tasks_enabled", "RetryPolicy", "TaskContext", "TaskMetadata", diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py index b4d7742527e0..b1588b4ef004 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -1,56 +1,58 @@ # --------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -"""Tests for opt-in gating of resilient ``TaskManager`` auto-initialization. - -``AgentServerHost`` must only stand up the ``TaskManager`` (and its -potentially network-backed startup recovery scan) when the application has -actually declared a durable task via ``@task`` / ``@multi_turn_task``. -Plain servers that never opt in — e.g. invocations-only hosts — must NOT -pay the hosted task-store startup cost (a blocking ``list()`` round-trip -plus credential-token acquisition that would gate server readiness while -having nothing to recover). - -The opt-in signal is a non-empty ``_REGISTERED_DESCRIPTORS`` registry. -Both durable decorators funnel through the single ``Task.__init__`` -registration site (``MultiTurnTask`` wraps an inner ``Task``), so declaring -either kind opts the app in. +"""Tests for the double-gated resilient ``TaskManager`` auto-initialization. + +``AgentServerHost`` stands up the resilient ``TaskManager`` (and its +potentially network-backed startup recovery scan) only when BOTH: + +1. the resilient task subsystem was explicitly enabled via + ``set_resilient_tasks_enabled(True)`` (defaults to ``False``), AND +2. at least one durable task has been declared (``@task`` / + ``@multi_turn_task``, tracked in the ``_REGISTERED_DESCRIPTORS`` list). + +Both conditions are read directly at lifespan startup. If either is false, +nothing is constructed and no task-store call is made — plain servers (e.g. +invocations-only hosts) pay nothing. """ import logging import pytest -from azure.ai.agentserver.core._base import _resilient_tasks_opted_in +from azure.ai.agentserver.core._base import ( + _has_registered_tasks, + _resilient_tasks_enabled, +) from azure.ai.agentserver.core.tasks import ( TaskContext, - TaskManagerNotInitialized, multi_turn_task, + resilient_tasks_enabled, + set_resilient_tasks_enabled, task, ) from azure.ai.agentserver.core.tasks import _decorator as _decorator_mod -from azure.ai.agentserver.core.tasks._manager import ( - get_task_manager, - set_task_manager, -) +from azure.ai.agentserver.core.tasks._manager import set_task_manager @pytest.fixture -def _isolate_registry(): - """Snapshot/clear/restore the global durable-task registry + manager. +def _clean_state(): + """Snapshot/clear/restore the global registry, manager, and enable switch. - ``_REGISTERED_DESCRIPTORS`` is process-global and populated at - decoration time by other test modules; snapshot and clear it so each - test controls the opt-in state, then restore it (and reset the manager - singleton) afterwards. + ``_REGISTERED_DESCRIPTORS`` and the enable switch are process-global and + may be touched by other test modules; snapshot and reset them so each test + controls the gate, then restore afterwards. """ - saved = list(_decorator_mod._REGISTERED_DESCRIPTORS) + saved_desc = list(_decorator_mod._REGISTERED_DESCRIPTORS) + saved_enabled = resilient_tasks_enabled() _decorator_mod._REGISTERED_DESCRIPTORS.clear() + set_resilient_tasks_enabled(False) set_task_manager(None) try: - yield _decorator_mod._REGISTERED_DESCRIPTORS + yield finally: _decorator_mod._REGISTERED_DESCRIPTORS.clear() - _decorator_mod._REGISTERED_DESCRIPTORS.extend(saved) + _decorator_mod._REGISTERED_DESCRIPTORS.extend(saved_desc) + set_resilient_tasks_enabled(saved_enabled) set_task_manager(None) @@ -63,12 +65,8 @@ def __init__(self, **kwargs) -> None: self.kwargs = kwargs self.startup_called = False self.shutdown_called = False - self.descriptors_loaded = False _FakeTaskManager.instances.append(self) - def load_registered_descriptors(self) -> None: - self.descriptors_loaded = True - async def startup(self) -> None: self.startup_called = True @@ -78,8 +76,7 @@ async def shutdown(self) -> None: @pytest.fixture def _fake_task_manager(monkeypatch: pytest.MonkeyPatch): - """Patch the ``TaskManager`` the lifespan imports so no real provider, - token acquisition, or task-store round-trip happens during the test.""" + """Patch ``TaskManager`` so no real provider/token/task-store call happens.""" _FakeTaskManager.instances.clear() monkeypatch.setattr( "azure.ai.agentserver.core.tasks._manager.TaskManager", @@ -88,129 +85,100 @@ def _fake_task_manager(monkeypatch: pytest.MonkeyPatch): return _FakeTaskManager +def _declare_task(name: str = "gate_probe") -> None: + @task(name=name) + async def _probe(ctx: "TaskContext[dict]") -> None: + return None + + # ------------------------------------------------------------------ # -# _resilient_tasks_opted_in(): the opt-in signal +# The two gate signals # ------------------------------------------------------------------ # -class TestOptInSignal: - """Unit tests for ``_resilient_tasks_opted_in()``.""" +class TestGateSignals: + """Unit tests for the enable switch and the descriptor-list check.""" - def test_empty_registry_is_not_opted_in(self, _isolate_registry) -> None: - assert _resilient_tasks_opted_in() is False + def test_switch_defaults_false(self, _clean_state) -> None: + assert resilient_tasks_enabled() is False + assert _resilient_tasks_enabled() is False - def test_task_decorator_opts_in(self, _isolate_registry) -> None: - @task(name="optin_probe_one_shot") - async def _probe(ctx: "TaskContext[dict]") -> None: - return None + def test_switch_toggles(self, _clean_state) -> None: + set_resilient_tasks_enabled(True) + assert _resilient_tasks_enabled() is True + set_resilient_tasks_enabled(False) + assert _resilient_tasks_enabled() is False + + def test_set_with_no_arg_enables(self, _clean_state) -> None: + set_resilient_tasks_enabled() + assert resilient_tasks_enabled() is True + + def test_has_registered_tasks_empty(self, _clean_state) -> None: + assert _has_registered_tasks() is False - assert _resilient_tasks_opted_in() is True + def test_has_registered_tasks_after_task(self, _clean_state) -> None: + _declare_task() + assert _has_registered_tasks() is True - def test_multi_turn_task_decorator_opts_in(self, _isolate_registry) -> None: - @multi_turn_task(name="optin_probe_multi_turn") + def test_has_registered_tasks_after_multi_turn(self, _clean_state) -> None: + @multi_turn_task(name="gate_probe_mt") async def _probe(ctx: "TaskContext[dict]") -> None: return None - assert _resilient_tasks_opted_in() is True - - def test_get_task_manager_still_raises_without_a_factory( - self, _isolate_registry - ) -> None: - """Outside an active host lifespan (no factory installed), the raise - behavior is preserved — lazy bootstrap only applies within a host.""" - with pytest.raises(TaskManagerNotInitialized): - get_task_manager() + assert _has_registered_tasks() is True # ------------------------------------------------------------------ # -# Lifespan behaviour: gate is honoured at startup +# The AND gate at lifespan startup # ------------------------------------------------------------------ # -class TestLifespanOptInGate: - """Verify ``AgentServerHost`` lifespan honours the opt-in gate.""" +class TestLifespanAndGate: + """The TaskManager runs iff (enabled AND at least one task declared).""" @pytest.mark.asyncio - async def test_not_opted_in_defers_eager_init( - self, - _isolate_registry, - _fake_task_manager, - caplog: pytest.LogCaptureFixture, - ) -> None: - """No durable task declared → no manager is eagerly constructed and - no recovery scan runs; the expensive startup path is deferred.""" + async def test_disabled_and_no_task_skips(self, _clean_state, _fake_task_manager) -> None: from azure.ai.agentserver.core import AgentServerHost app = AgentServerHost() + async with app.router.lifespan_context(app): + pass + assert _fake_task_manager.instances == [] - with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): - async with app.router.lifespan_context(app): - # Nothing eagerly built during startup — no provider / token / - # task-store round-trip happened. - assert _fake_task_manager.instances == [] + @pytest.mark.asyncio + async def test_disabled_but_task_declared_skips(self, _clean_state, _fake_task_manager) -> None: + """Flag off wins: a declared task alone does NOT start the manager.""" + from azure.ai.agentserver.core import AgentServerHost - assert not any( - "TaskManager initialized automatically" in r.message for r in caplog.records - ) - assert any( - "TaskManager auto-init deferred" in r.message for r in caplog.records - ) + _declare_task() + # switch left at default False + app = AgentServerHost() + async with app.router.lifespan_context(app): + pass + assert _fake_task_manager.instances == [] @pytest.mark.asyncio - async def test_first_task_declared_in_lifespan_lazily_bootstraps( - self, - _isolate_registry, - _fake_task_manager, - ) -> None: - """Late-registration path: a task declared *during* the active - lifespan (when eager init was skipped) must not fail on first use — - ``get_task_manager()`` lazily bootstraps a manager instead of raising - ``TaskManagerNotInitialized``.""" + async def test_enabled_but_no_task_skips(self, _clean_state, _fake_task_manager) -> None: + """Switch on but no task declared: registry empty -> manager not built.""" from azure.ai.agentserver.core import AgentServerHost + set_resilient_tasks_enabled(True) app = AgentServerHost() - async with app.router.lifespan_context(app): - # Nothing built yet. - assert _fake_task_manager.instances == [] - - # Declare the FIRST durable task now, mid-lifespan. - @task(name="optin_late_declared_task") - async def _late(ctx: "TaskContext[dict]") -> None: - return None - - # First use resolves the manager — must not raise, and must - # lazily construct + load descriptors. This is exactly the first - # thing ``Task.run``/``Task.start`` do, so it protects the - # advertised "first run/start no longer raises" behavior. - mgr = get_task_manager() - assert mgr is _fake_task_manager.instances[0] - assert mgr.descriptors_loaded is True - # Intentional limitation: the sync lazy path does NOT run the - # async ``startup()`` (initial recovery scan + periodic loop). - # A task declared after startup already missed the startup - # recovery window; its own run/start works via the provider. - assert mgr.startup_called is False - - # The lazily-bootstrapped manager is torn down on shutdown. - assert _fake_task_manager.instances[0].shutdown_called is True + pass + assert _fake_task_manager.instances == [] @pytest.mark.asyncio - async def test_opted_in_initializes_task_manager( - self, - _isolate_registry, - _fake_task_manager, - caplog: pytest.LogCaptureFixture, + async def test_enabled_and_task_initializes( + self, _clean_state, _fake_task_manager, caplog: pytest.LogCaptureFixture ) -> None: - """A declared ``@task`` → TaskManager is constructed and started.""" + """Both conditions true -> manager constructed, started, shut down.""" from azure.ai.agentserver.core import AgentServerHost - @task(name="optin_lifespan_task") - async def _probe(ctx: "TaskContext[dict]") -> None: - return None + set_resilient_tasks_enabled(True) + _declare_task() app = AgentServerHost() - with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): async with app.router.lifespan_context(app): pass @@ -219,55 +187,21 @@ async def _probe(ctx: "TaskContext[dict]") -> None: mgr = _fake_task_manager.instances[0] assert mgr.startup_called is True assert mgr.shutdown_called is True - assert any( - "TaskManager initialized automatically" in r.message for r in caplog.records - ) + assert any("TaskManager initialized automatically" in r.message for r in caplog.records) @pytest.mark.asyncio - async def test_multi_turn_opt_in_initializes_task_manager( - self, - _isolate_registry, - _fake_task_manager, - ) -> None: - """A declared ``@multi_turn_task`` also opts the host in.""" + async def test_enabled_and_multi_turn_initializes(self, _clean_state, _fake_task_manager) -> None: from azure.ai.agentserver.core import AgentServerHost - @multi_turn_task(name="optin_lifespan_multi_turn") + set_resilient_tasks_enabled(True) + + @multi_turn_task(name="gate_lifespan_mt") async def _probe(ctx: "TaskContext[dict]") -> None: return None app = AgentServerHost() - async with app.router.lifespan_context(app): pass assert len(_fake_task_manager.instances) == 1 assert _fake_task_manager.instances[0].startup_called is True - - @pytest.mark.asyncio - async def test_lifespan_tears_down_cleanly_when_tasks_module_unavailable( - self, - _isolate_registry, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - """Task-less / optional install: both startup AND shutdown tolerate - the resilient-task module being unimportable (no crash on teardown).""" - import builtins - - real_import = builtins.__import__ - - def _blocked_import(name, *args, **kwargs): - if name.endswith("tasks._manager") or name == "azure.ai.agentserver.core.tasks._manager": - raise ImportError("simulated: resilient tasks module unavailable") - return real_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, "__import__", _blocked_import) - - from azure.ai.agentserver.core import AgentServerHost - - app = AgentServerHost() - - # Entering AND exiting the lifespan must both succeed without raising - # ImportError — the shutdown import is guarded like the startup one. - async with app.router.lifespan_context(app): - pass From 951c6a316805b1b8cb93c357e2fc8d679c35dd34 Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Fri, 31 Jul 2026 01:05:28 +0530 Subject: [PATCH 5/7] Always construct TaskManager; gate only the network recovery scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refine the design so it fixes the startup latency WITHOUT the backward-compat break the reviewer (Copilot) flagged: existing @task apps that don't call the new switch previously would hit TaskManagerNotInitialized on first run/start. Now the TaskManager is CONSTRUCTED unconditionally at lifespan startup — construction is cheap and makes no task-store calls (only in-memory state). This keeps get_task_manager()/.run()/.start() working for existing task apps. The network-backed startup RECOVERY SCAN (the blocking hosted task-store list() + credential-token acquisition that caused the 20-36s stall) runs only when BOTH set_resilient_tasks_enabled(True) was called AND at least one durable task is declared. Plain servers make no task-store call at startup. - _base.py: construct + set_task_manager always; gate await startup() on (enabled AND registered). shutdown() is safe on a never-started manager. - Tests rewritten: manager always constructed + retrievable via get_task_manager(); startup_called only in the enabled+task case. - Fix second public-API guard set (test_contract_completeness) that CI caught; updated tasks-guide.md + CHANGELOG to the always-construct semantics. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5222ab61-7d5e-4642-b75f-ef365e375f4f --- .../azure-ai-agentserver-core/CHANGELOG.md | 4 +- .../azure/ai/agentserver/core/_base.py | 76 ++++++++++--------- .../docs/tasks-guide.md | 16 ++-- .../tests/tasks/test_contract_completeness.py | 3 + .../tests/tasks/test_task_manager_optin.py | 46 +++++++---- 5 files changed, 87 insertions(+), 58 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 71407e8058cb..dd0857151f24 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -5,11 +5,11 @@ ### Features Added - Added public `MiddlewareFactory` and `StreamContent` typing aliases for host middleware and streaming helpers. -- Added `set_resilient_tasks_enabled` / `resilient_tasks_enabled` to `azure.ai.agentserver.core.tasks` — an explicit process-global switch (default **disabled**) that gates the resilient `TaskManager`. +- Added `set_resilient_tasks_enabled` / `resilient_tasks_enabled` to `azure.ai.agentserver.core.tasks` — an explicit process-global switch (default **disabled**) that gates the resilient `TaskManager`'s startup recovery scan. ### Bugs Fixed -- `AgentServerHost` now initializes the resilient `TaskManager` (and its network-backed startup recovery scan) only when **both** the subsystem was explicitly enabled via `set_resilient_tasks_enabled(True)` **and** at least one durable task (`@task` / `@multi_turn_task`) has been declared. Plain servers (e.g. invocations-only hosts) that do not opt in no longer pay the task-store startup cost — a blocking hosted task-store `list()` round-trip plus credential-token acquisition that could gate server readiness and add tens of seconds of connect latency while having nothing to recover. +- `AgentServerHost` no longer makes a blocking hosted task-store `list()` round-trip (plus credential-token acquisition) at startup unless resilient tasks are actually in use. The `TaskManager` is still constructed so `get_task_manager()` keeps working (no `TaskManagerNotInitialized` for existing `@task` apps), but its network-backed **startup recovery scan** now runs only when **both** the subsystem was explicitly enabled via `set_resilient_tasks_enabled(True)` **and** at least one durable task (`@task` / `@multi_turn_task`) has been declared. Plain servers (e.g. invocations-only hosts) no longer pay tens of seconds of startup latency for a scan that has nothing to recover. - Cleaned up `AgentServerHost` public signatures so inherited middleware typing does not expose Starlette private type aliases. ## 2.0.0b9 (2026-07-28) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index ea785be98855..dbee24ecc761 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -311,47 +311,53 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF protocols, ) - # --- Resilient task manager auto-initialization --- + # --- Resilient task manager initialization --- # - # DOUBLE GATE (AND): the TaskManager — and its network-backed - # startup recovery scan (a blocking hosted task-store ``list()`` - # plus DefaultAzureCredential token acquisition) — is stood up ONLY - # when BOTH hold: + # The TaskManager is CONSTRUCTED unconditionally (whenever the + # resilient tasks module is importable). Construction is cheap and + # makes NO task-store calls — it only builds in-memory state (an + # idle provider client, empty routing tables, a lease-owner string). + # This keeps ``get_task_manager()`` working so a ``@task``-based app + # can run tasks without hitting ``TaskManagerNotInitialized``, + # while paying zero network cost until a task is actually used. + # + # The network-backed startup RECOVERY SCAN (a blocking hosted + # task-store ``list()`` plus ``DefaultAzureCredential`` token + # acquisition, which would otherwise gate server readiness) runs + # ONLY when BOTH hold: # (1) resilient tasks were explicitly enabled via # ``set_resilient_tasks_enabled(True)`` (default False), AND - # (2) the ``_REGISTERED_DESCRIPTORS`` list is non-empty, i.e. at - # least one durable task (``@task`` / ``@multi_turn_task``) - # has been declared. - # - # Both are read directly here. If either is false, nothing is - # constructed and no task-store call is made — plain servers (e.g. - # invocations-only hosts) pay nothing. + # (2) the ``_REGISTERED_DESCRIPTORS`` list is non-empty (at least + # one ``@task`` / ``@multi_turn_task`` was declared). + # So plain servers (e.g. invocations-only hosts) make no task-store + # call at startup, and the switch gates all eager task-store work. task_manager = None - if _resilient_tasks_enabled() and _has_registered_tasks(): - try: - from .tasks._manager import ( # pylint: disable=import-outside-toplevel - TaskManager, - set_task_manager, - ) + try: + from .tasks._manager import ( # pylint: disable=import-outside-toplevel + TaskManager, + set_task_manager, + ) - task_manager = TaskManager( - config=cfg, - shutdown_event=asyncio.Event(), - shutdown_grace_seconds=_read_task_manager_shutdown_grace(), - ) - set_task_manager(task_manager) - await task_manager.startup() - logger.info("TaskManager initialized automatically") - except ImportError: - pass # resilient module not available - except Exception: # pylint: disable=broad-exception-caught - logger.warning("Failed to initialize TaskManager", exc_info=True) - else: - logger.info( - "TaskManager not initialized (enabled=%s, tasks_declared=%s)", - _resilient_tasks_enabled(), - _has_registered_tasks(), + task_manager = TaskManager( + config=cfg, + shutdown_event=asyncio.Event(), + shutdown_grace_seconds=_read_task_manager_shutdown_grace(), ) + set_task_manager(task_manager) + + if _resilient_tasks_enabled() and _has_registered_tasks(): + await task_manager.startup() + logger.info("TaskManager initialized with startup recovery") + else: + logger.info( + "TaskManager initialized (recovery deferred; enabled=%s, tasks_declared=%s)", + _resilient_tasks_enabled(), + _has_registered_tasks(), + ) + except ImportError: + pass # resilient module not available + except Exception: # pylint: disable=broad-exception-caught + logger.warning("Failed to initialize TaskManager", exc_info=True) yield diff --git a/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md b/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md index 0a8f99d3574e..6be906ad6d55 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md @@ -115,9 +115,10 @@ What this primitive deliberately does **not** do: ### Enabling resilient tasks -The resilient `TaskManager` — and its startup recovery scan against the hosted -task store — is **opt-in**. `AgentServerHost` only stands it up when **both** -of the following are true: +The resilient `TaskManager`'s **startup recovery scan** — a network round-trip +to the hosted task store that reclaims tasks left in-flight by a crashed prior +instance — is **opt-in**. `AgentServerHost` runs it only when **both** of the +following are true: 1. the subsystem was explicitly enabled via `set_resilient_tasks_enabled(True)` (it defaults to disabled), **and** @@ -129,11 +130,14 @@ module import). Use `resilient_tasks_enabled()` to read the current state. ```python from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled -set_resilient_tasks_enabled(True) # required for the TaskManager to run +set_resilient_tasks_enabled(True) # run the startup recovery scan for tasks ``` -Servers that never declare a task — or never enable the switch — skip the -TaskManager entirely and pay none of its startup cost. +The `TaskManager` itself is always constructed (a cheap, in-memory object that +makes no task-store calls until a task is used), so `get_task_manager()` and +`.run()` / `.start()` work regardless. Servers that never declare a task — or +never enable the switch — simply skip the startup recovery scan and pay none of +its latency. ### One-shot diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_contract_completeness.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_contract_completeness.py index e3a65c284df4..89785185ee80 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_contract_completeness.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_contract_completeness.py @@ -79,6 +79,9 @@ "multi_turn_task", "Task", "MultiTurnTask", + # Enablement switch + "set_resilient_tasks_enabled", + "resilient_tasks_enabled", # Context + metadata "TaskContext", "TaskMetadata", diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py index b1588b4ef004..421a10c93e5a 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -133,21 +133,35 @@ async def _probe(ctx: "TaskContext[dict]") -> None: # ------------------------------------------------------------------ # -class TestLifespanAndGate: - """The TaskManager runs iff (enabled AND at least one task declared).""" +class TestLifespanManagerAndRecovery: + """The TaskManager is ALWAYS constructed (cheap, no task-store calls); + the network-backed startup recovery scan runs only when BOTH the switch + is enabled AND at least one durable task is declared.""" @pytest.mark.asyncio - async def test_disabled_and_no_task_skips(self, _clean_state, _fake_task_manager) -> None: + async def test_manager_always_constructed_and_available( + self, _clean_state, _fake_task_manager + ) -> None: + """Even with no opt-in, the manager is constructed and installed so + ``get_task_manager()`` works (no ``TaskManagerNotInitialized``) — but + its startup recovery scan does NOT run.""" from azure.ai.agentserver.core import AgentServerHost + from azure.ai.agentserver.core.tasks._manager import get_task_manager app = AgentServerHost() async with app.router.lifespan_context(app): - pass - assert _fake_task_manager.instances == [] + # A manager exists and is retrievable during the active lifespan. + assert len(_fake_task_manager.instances) == 1 + assert get_task_manager() is _fake_task_manager.instances[0] + # No recovery scan happened (no opt-in). + assert _fake_task_manager.instances[0].startup_called is False + + # Torn down + cleared on shutdown. + assert _fake_task_manager.instances[0].shutdown_called is True @pytest.mark.asyncio - async def test_disabled_but_task_declared_skips(self, _clean_state, _fake_task_manager) -> None: - """Flag off wins: a declared task alone does NOT start the manager.""" + async def test_disabled_but_task_declared_no_recovery(self, _clean_state, _fake_task_manager) -> None: + """Switch off + a task declared: manager built, but no recovery scan.""" from azure.ai.agentserver.core import AgentServerHost _declare_task() @@ -155,24 +169,26 @@ async def test_disabled_but_task_declared_skips(self, _clean_state, _fake_task_m app = AgentServerHost() async with app.router.lifespan_context(app): pass - assert _fake_task_manager.instances == [] + assert len(_fake_task_manager.instances) == 1 + assert _fake_task_manager.instances[0].startup_called is False @pytest.mark.asyncio - async def test_enabled_but_no_task_skips(self, _clean_state, _fake_task_manager) -> None: - """Switch on but no task declared: registry empty -> manager not built.""" + async def test_enabled_but_no_task_no_recovery(self, _clean_state, _fake_task_manager) -> None: + """Switch on + no task declared: manager built, but no recovery scan.""" from azure.ai.agentserver.core import AgentServerHost set_resilient_tasks_enabled(True) app = AgentServerHost() async with app.router.lifespan_context(app): pass - assert _fake_task_manager.instances == [] + assert len(_fake_task_manager.instances) == 1 + assert _fake_task_manager.instances[0].startup_called is False @pytest.mark.asyncio - async def test_enabled_and_task_initializes( + async def test_enabled_and_task_runs_recovery( self, _clean_state, _fake_task_manager, caplog: pytest.LogCaptureFixture ) -> None: - """Both conditions true -> manager constructed, started, shut down.""" + """Both conditions true -> manager built AND startup recovery runs.""" from azure.ai.agentserver.core import AgentServerHost set_resilient_tasks_enabled(True) @@ -187,10 +203,10 @@ async def test_enabled_and_task_initializes( mgr = _fake_task_manager.instances[0] assert mgr.startup_called is True assert mgr.shutdown_called is True - assert any("TaskManager initialized automatically" in r.message for r in caplog.records) + assert any("TaskManager initialized with startup recovery" in r.message for r in caplog.records) @pytest.mark.asyncio - async def test_enabled_and_multi_turn_initializes(self, _clean_state, _fake_task_manager) -> None: + async def test_enabled_and_multi_turn_runs_recovery(self, _clean_state, _fake_task_manager) -> None: from azure.ai.agentserver.core import AgentServerHost set_resilient_tasks_enabled(True) From f3da0cf667755b79cee5db1921a13c4a2f899956 Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Fri, 31 Jul 2026 01:50:58 +0530 Subject: [PATCH 6/7] Gate recovery scan on (task declared OR switch) instead of AND MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch the double gate from AND to OR so an existing @task app keeps automatic crash recovery WITHOUT calling the new switch (backward compatible), and the switch becomes a force-enable that starts the recovery loop even before any task is declared — which also closes the recovery gap for tasks registered lazily under an enabled switch (the running periodic loop picks them up dynamically), without touching _manager.py. Recovery scan now runs when: at least one durable task is declared (@task / @multi_turn_task) OR set_resilient_tasks_enabled(True) was set. The TaskManager is still always constructed (no TaskManagerNotInitialized). Deferred: if the switch is off AND no task is declared at startup, a task declared later that lifetime runs but its prior-crash orphans aren't scanned until the next restart (no recovery loop was started). Fully closing this needs a lazy manager-start on first late registration — tracked as future work (the manager can be made lazy, not fully async). Also address review: reword set_resilient_tasks_enabled docstring (it force-enables recovery, does not gate .run()/.start()); add the switch to the guide's standalone examples; add release date; drop the breaking-change framing (OR is backward compatible). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5222ab61-7d5e-4642-b75f-ef365e375f4f --- .../azure-ai-agentserver-core/CHANGELOG.md | 6 +-- .../azure/ai/agentserver/core/_base.py | 30 +++++++++---- .../ai/agentserver/core/tasks/_enablement.py | 45 ++++++++++++------- .../docs/tasks-guide.md | 36 +++++++++------ .../tests/tasks/test_task_manager_optin.py | 35 ++++++++------- 5 files changed, 94 insertions(+), 58 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index dd0857151f24..0327b674ca21 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -1,15 +1,15 @@ # Release History -## 2.0.0b10 (Unreleased) +## 2.0.0b10 (2026-07-31) ### Features Added - Added public `MiddlewareFactory` and `StreamContent` typing aliases for host middleware and streaming helpers. -- Added `set_resilient_tasks_enabled` / `resilient_tasks_enabled` to `azure.ai.agentserver.core.tasks` — an explicit process-global switch (default **disabled**) that gates the resilient `TaskManager`'s startup recovery scan. +- Added `set_resilient_tasks_enabled` / `resilient_tasks_enabled` to `azure.ai.agentserver.core.tasks` — a process-global switch (default off) that force-enables the resilient `TaskManager`'s startup recovery scan even before any durable task is declared (useful when tasks are registered lazily after startup). ### Bugs Fixed -- `AgentServerHost` no longer makes a blocking hosted task-store `list()` round-trip (plus credential-token acquisition) at startup unless resilient tasks are actually in use. The `TaskManager` is still constructed so `get_task_manager()` keeps working (no `TaskManagerNotInitialized` for existing `@task` apps), but its network-backed **startup recovery scan** now runs only when **both** the subsystem was explicitly enabled via `set_resilient_tasks_enabled(True)` **and** at least one durable task (`@task` / `@multi_turn_task`) has been declared. Plain servers (e.g. invocations-only hosts) no longer pay tens of seconds of startup latency for a scan that has nothing to recover. +- `AgentServerHost` no longer makes a blocking hosted task-store `list()` round-trip (plus credential-token acquisition) at startup unless resilient tasks are actually in use. The `TaskManager` is still constructed so `get_task_manager()` and `.run()` / `.start()` keep working, but its network-backed **startup recovery scan** now runs only when **either** at least one durable task (`@task` / `@multi_turn_task`) has been declared **or** `set_resilient_tasks_enabled(True)` was called. Apps that use tasks keep automatic crash recovery with no code change; plain servers (e.g. invocations-only hosts) no longer pay tens of seconds of startup latency for a scan that has nothing to recover. - Cleaned up `AgentServerHost` public signatures so inherited middleware typing does not expose Starlette private type aliases. ## 2.0.0b9 (2026-07-28) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index dbee24ecc761..1082467956ae 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -323,14 +323,26 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF # # The network-backed startup RECOVERY SCAN (a blocking hosted # task-store ``list()`` plus ``DefaultAzureCredential`` token - # acquisition, which would otherwise gate server readiness) runs - # ONLY when BOTH hold: - # (1) resilient tasks were explicitly enabled via - # ``set_resilient_tasks_enabled(True)`` (default False), AND - # (2) the ``_REGISTERED_DESCRIPTORS`` list is non-empty (at least - # one ``@task`` / ``@multi_turn_task`` was declared). - # So plain servers (e.g. invocations-only hosts) make no task-store - # call at startup, and the switch gates all eager task-store work. + # acquisition, which would otherwise gate server readiness) — and + # the periodic recovery loop it spawns — run when EITHER holds: + # (1) at least one durable task was declared (``@task`` / + # ``@multi_turn_task``, tracked in ``_REGISTERED_DESCRIPTORS``) + # — an app that uses tasks gets recovery automatically, OR + # (2) resilient tasks were explicitly enabled via + # ``set_resilient_tasks_enabled(True)`` — a force-enable that + # starts the recovery loop even before any task is declared, + # so a task declared later is picked up by the loop. + # A plain server that neither declares a task nor sets the switch + # (e.g. an invocations-only host) makes no task-store call at + # startup. + # + # NOTE (deferred): if the switch is OFF and no task is declared at + # startup, the recovery loop is not started, so a task declared + # LATER in that lifetime will run but its prior-crash orphans are + # not scanned until the next restart. Fully closing that requires a + # lazy manager-start on first late registration; tracked as future + # work (the manager cannot be made fully async, only lazily + # started). task_manager = None try: from .tasks._manager import ( # pylint: disable=import-outside-toplevel @@ -345,7 +357,7 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF ) set_task_manager(task_manager) - if _resilient_tasks_enabled() and _has_registered_tasks(): + if _resilient_tasks_enabled() or _has_registered_tasks(): await task_manager.startup() logger.info("TaskManager initialized with startup recovery") else: diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py index eadc355b9f4e..892ef5e9f41f 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py @@ -1,20 +1,24 @@ # --------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -"""Explicit enablement switch for the resilient task subsystem. +"""Explicit force-enable switch for the resilient task recovery scan. -The resilient ``TaskManager`` performs network-backed startup work (a hosted -task-store recovery scan plus credential-token acquisition). To avoid making -any of those calls unless an app genuinely relies on durable tasks, -``AgentServerHost`` only stands up the ``TaskManager`` when BOTH conditions -hold: +The resilient ``TaskManager`` is always constructed by ``AgentServerHost`` (a +cheap, in-memory object that makes no task-store calls), so ``get_task_manager`` +and ``.run()`` / ``.start()`` work regardless of this switch. What this switch +affects is only the network-backed **startup recovery scan** (a hosted +task-store ``list()`` plus credential-token acquisition) and the periodic +recovery loop it spawns. -1. this switch has been turned on via :func:`set_resilient_tasks_enabled`, AND -2. at least one durable task has been declared (``@task`` / ``@multi_turn_task``). +That recovery runs at startup when EITHER a durable task has been declared +(``@task`` / ``@multi_turn_task``) OR this switch is on. So an app that uses +tasks gets recovery automatically; this switch is a **force-enable** that +starts the recovery loop even before any task is declared (useful when tasks +are registered lazily after startup — the running loop then picks them up). -The switch is process-global and defaults to ``False`` (disabled). It is -intentionally decoupled from any ``AgentServerHost`` instance so it can be -flipped independently at import time, e.g.:: +The switch is process-global and defaults to ``False``. It is intentionally +decoupled from any ``AgentServerHost`` instance so it can be flipped +independently at import time, e.g.:: from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled @@ -25,13 +29,18 @@ def set_resilient_tasks_enabled(value: bool = True) -> None: - """Enable or disable the resilient task subsystem process-wide. + """Force-enable (or clear) the resilient task recovery scan process-wide. + + Setting this to ``True`` starts the startup recovery scan + periodic + recovery loop even when no durable task is declared at startup. It does + NOT gate the ``TaskManager``'s existence: ``.run()`` / ``.start()`` work + whether or not this is set — this only controls automatic crash recovery. Must be called before ``AgentServerHost`` lifespan startup (typically at - import time) to take effect for eager initialization. Defaults the switch - to enabled when called with no argument. + import time) to take effect. Defaults to enabling when called with no + argument. - :param value: ``True`` to enable resilient tasks, ``False`` to disable. + :param value: ``True`` to force-enable recovery, ``False`` to clear. :type value: bool """ global _RESILIENT_TASKS_ENABLED # pylint: disable=global-statement @@ -39,7 +48,11 @@ def set_resilient_tasks_enabled(value: bool = True) -> None: def resilient_tasks_enabled() -> bool: - """Return whether the resilient task subsystem has been explicitly enabled. + """Return whether the recovery scan was explicitly force-enabled. + + Note this reflects only the switch — recovery also runs automatically when + a durable task is declared, so a ``False`` return does not mean recovery is + off, nor that tasks are unavailable. :return: ``True`` if :func:`set_resilient_tasks_enabled` turned it on. :rtype: bool diff --git a/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md b/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md index 6be906ad6d55..c8103665f70e 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md @@ -117,33 +117,39 @@ What this primitive deliberately does **not** do: The resilient `TaskManager`'s **startup recovery scan** — a network round-trip to the hosted task store that reclaims tasks left in-flight by a crashed prior -instance — is **opt-in**. `AgentServerHost` runs it only when **both** of the -following are true: +instance — runs at startup when **either** of the following holds: -1. the subsystem was explicitly enabled via `set_resilient_tasks_enabled(True)` - (it defaults to disabled), **and** -2. at least one durable task has been declared (`@task` / `@multi_turn_task`). +1. at least one durable task has been declared (`@task` / `@multi_turn_task`) — + an app that uses tasks gets recovery automatically, **or** +2. it was explicitly force-enabled via `set_resilient_tasks_enabled(True)`. -Call `set_resilient_tasks_enabled(True)` once, before the host starts (e.g. at -module import). Use `resilient_tasks_enabled()` to read the current state. +The force-enable is useful when tasks are registered *lazily* (declared after +startup): it starts the periodic recovery loop up front so a task declared +later is still recovered. ```python from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled -set_resilient_tasks_enabled(True) # run the startup recovery scan for tasks +set_resilient_tasks_enabled(True) # force-enable recovery before any task ``` +Use `resilient_tasks_enabled()` to read the current switch state. + The `TaskManager` itself is always constructed (a cheap, in-memory object that makes no task-store calls until a task is used), so `get_task_manager()` and -`.run()` / `.start()` work regardless. Servers that never declare a task — or -never enable the switch — simply skip the startup recovery scan and pay none of -its latency. +`.run()` / `.start()` work regardless of the switch. A server that neither +declares a task nor sets the switch (e.g. an invocations-only host) simply +skips the startup recovery scan and pays none of its latency. ### One-shot ```python import asyncio -from azure.ai.agentserver.core.tasks import task, TaskContext +from azure.ai.agentserver.core.tasks import task, TaskContext, set_resilient_tasks_enabled + +# Optional: force-enable recovery even before the first task is declared. +# Declaring the @task below already enables recovery on its own. +set_resilient_tasks_enabled(True) @task(name="summarize") async def summarize(ctx: TaskContext[str]) -> str: @@ -163,7 +169,11 @@ asyncio.run(main()) ```python import asyncio -from azure.ai.agentserver.core.tasks import multi_turn_task, TaskContext +from azure.ai.agentserver.core.tasks import multi_turn_task, TaskContext, set_resilient_tasks_enabled + +# Optional: force-enable recovery even before the first task is declared. +# Declaring the @multi_turn_task below already enables recovery on its own. +set_resilient_tasks_enabled(True) @multi_turn_task(name="chat") async def chat(ctx: TaskContext[dict]) -> dict: diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py index 421a10c93e5a..15f85605a6b1 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -135,16 +135,16 @@ async def _probe(ctx: "TaskContext[dict]") -> None: class TestLifespanManagerAndRecovery: """The TaskManager is ALWAYS constructed (cheap, no task-store calls); - the network-backed startup recovery scan runs only when BOTH the switch - is enabled AND at least one durable task is declared.""" + the network-backed startup recovery scan runs when EITHER the switch is + enabled OR at least one durable task is declared.""" @pytest.mark.asyncio - async def test_manager_always_constructed_and_available( + async def test_neither_enabled_nor_task_no_recovery( self, _clean_state, _fake_task_manager ) -> None: - """Even with no opt-in, the manager is constructed and installed so + """No switch, no task: the manager is constructed and installed so ``get_task_manager()`` works (no ``TaskManagerNotInitialized``) — but - its startup recovery scan does NOT run.""" + the startup recovery scan does NOT run (plain invocations host).""" from azure.ai.agentserver.core import AgentServerHost from azure.ai.agentserver.core.tasks._manager import get_task_manager @@ -153,15 +153,16 @@ async def test_manager_always_constructed_and_available( # A manager exists and is retrievable during the active lifespan. assert len(_fake_task_manager.instances) == 1 assert get_task_manager() is _fake_task_manager.instances[0] - # No recovery scan happened (no opt-in). + # No recovery scan happened (neither gate true). assert _fake_task_manager.instances[0].startup_called is False # Torn down + cleared on shutdown. assert _fake_task_manager.instances[0].shutdown_called is True @pytest.mark.asyncio - async def test_disabled_but_task_declared_no_recovery(self, _clean_state, _fake_task_manager) -> None: - """Switch off + a task declared: manager built, but no recovery scan.""" + async def test_task_declared_runs_recovery_without_switch(self, _clean_state, _fake_task_manager) -> None: + """A declared task alone runs recovery (backward compatible — an + existing ``@task`` app gets recovery without calling the switch).""" from azure.ai.agentserver.core import AgentServerHost _declare_task() @@ -170,11 +171,12 @@ async def test_disabled_but_task_declared_no_recovery(self, _clean_state, _fake_ async with app.router.lifespan_context(app): pass assert len(_fake_task_manager.instances) == 1 - assert _fake_task_manager.instances[0].startup_called is False + assert _fake_task_manager.instances[0].startup_called is True @pytest.mark.asyncio - async def test_enabled_but_no_task_no_recovery(self, _clean_state, _fake_task_manager) -> None: - """Switch on + no task declared: manager built, but no recovery scan.""" + async def test_switch_alone_runs_recovery_without_task(self, _clean_state, _fake_task_manager) -> None: + """The switch alone runs recovery (force-enable) — starting the + periodic recovery loop so a task declared later is picked up.""" from azure.ai.agentserver.core import AgentServerHost set_resilient_tasks_enabled(True) @@ -182,13 +184,13 @@ async def test_enabled_but_no_task_no_recovery(self, _clean_state, _fake_task_ma async with app.router.lifespan_context(app): pass assert len(_fake_task_manager.instances) == 1 - assert _fake_task_manager.instances[0].startup_called is False + assert _fake_task_manager.instances[0].startup_called is True @pytest.mark.asyncio - async def test_enabled_and_task_runs_recovery( + async def test_switch_and_task_runs_recovery( self, _clean_state, _fake_task_manager, caplog: pytest.LogCaptureFixture ) -> None: - """Both conditions true -> manager built AND startup recovery runs.""" + """Both true -> manager built AND startup recovery runs.""" from azure.ai.agentserver.core import AgentServerHost set_resilient_tasks_enabled(True) @@ -206,11 +208,10 @@ async def test_enabled_and_task_runs_recovery( assert any("TaskManager initialized with startup recovery" in r.message for r in caplog.records) @pytest.mark.asyncio - async def test_enabled_and_multi_turn_runs_recovery(self, _clean_state, _fake_task_manager) -> None: + async def test_multi_turn_task_runs_recovery_without_switch(self, _clean_state, _fake_task_manager) -> None: + """A declared ``@multi_turn_task`` alone also runs recovery.""" from azure.ai.agentserver.core import AgentServerHost - set_resilient_tasks_enabled(True) - @multi_turn_task(name="gate_lifespan_mt") async def _probe(ctx: "TaskContext[dict]") -> None: return None From 2e390028554c3ccdc086807e6fbe329fbcd272e2 Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Fri, 31 Jul 2026 02:04:33 +0530 Subject: [PATCH 7/7] Fix stale AND wording in test module docstring (semantics are OR) The test module docstring still described the earlier AND gate; update it to the OR contract (recovery runs when a task is declared OR the switch is set) to match the implementation, tests, docs, and CHANGELOG. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5222ab61-7d5e-4642-b75f-ef365e375f4f --- .../tests/tasks/test_task_manager_optin.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py index 15f85605a6b1..6cff55e7d56f 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -1,19 +1,21 @@ # --------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -"""Tests for the double-gated resilient ``TaskManager`` auto-initialization. +"""Tests for the gated resilient ``TaskManager`` startup recovery scan. -``AgentServerHost`` stands up the resilient ``TaskManager`` (and its -potentially network-backed startup recovery scan) only when BOTH: +``AgentServerHost`` always constructs the resilient ``TaskManager`` (a cheap, +in-memory object that makes no task-store calls), so ``get_task_manager()`` and +``.run()`` / ``.start()`` work regardless. Its network-backed **startup +recovery scan** (and the periodic recovery loop it spawns) runs when EITHER: -1. the resilient task subsystem was explicitly enabled via - ``set_resilient_tasks_enabled(True)`` (defaults to ``False``), AND -2. at least one durable task has been declared (``@task`` / - ``@multi_turn_task``, tracked in the ``_REGISTERED_DESCRIPTORS`` list). +1. at least one durable task has been declared (``@task`` / + ``@multi_turn_task``, tracked in the ``_REGISTERED_DESCRIPTORS`` list), OR +2. the switch was set via ``set_resilient_tasks_enabled(True)`` (default + ``False``) — a force-enable. -Both conditions are read directly at lifespan startup. If either is false, -nothing is constructed and no task-store call is made — plain servers (e.g. -invocations-only hosts) pay nothing. +Both signals are read directly at lifespan startup. When neither is true, no +task-store call is made — plain servers (e.g. invocations-only hosts) pay +nothing. """ import logging @@ -190,7 +192,7 @@ async def test_switch_alone_runs_recovery_without_task(self, _clean_state, _fake async def test_switch_and_task_runs_recovery( self, _clean_state, _fake_task_manager, caplog: pytest.LogCaptureFixture ) -> None: - """Both true -> manager built AND startup recovery runs.""" + """Both signals true -> manager built and startup recovery runs.""" from azure.ai.agentserver.core import AgentServerHost set_resilient_tasks_enabled(True)