diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 50110d7b7c01..0327b674ca21 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -1,13 +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` — 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()` 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/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 e6efee6306ef..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 @@ -59,6 +59,51 @@ def _read_task_manager_shutdown_grace() -> float: return 25.0 +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: + from .tasks._decorator import ( # pylint: disable=import-outside-toplevel + _REGISTERED_DESCRIPTORS, + ) + except ImportError: + return False + 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. @@ -266,7 +311,38 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF protocols, ) - # --- Resilient task manager auto-initialization --- + # --- Resilient task manager initialization --- + # + # 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) — 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 @@ -280,8 +356,16 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF shutdown_grace_seconds=_read_task_manager_shutdown_grace(), ) set_task_manager(task_manager) - await task_manager.startup() - logger.info("TaskManager initialized automatically") + + if _resilient_tasks_enabled() or _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 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..892ef5e9f41f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_enablement.py @@ -0,0 +1,60 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Explicit force-enable switch for the resilient task recovery scan. + +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. + +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``. 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: + """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. Defaults to enabling when called with no + argument. + + :param value: ``True`` to force-enable recovery, ``False`` to clear. + :type value: bool + """ + global _RESILIENT_TASKS_ENABLED # pylint: disable=global-statement + _RESILIENT_TASKS_ENABLED = bool(value) + + +def resilient_tasks_enabled() -> bool: + """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 + """ + return _RESILIENT_TASKS_ENABLED 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..c8103665f70e 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/tasks-guide.md @@ -113,11 +113,43 @@ What this primitive deliberately does **not** do: ## 3. Hello world +### Enabling resilient tasks + +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 — runs at startup when **either** of the following holds: + +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)`. + +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) # 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 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: @@ -137,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_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_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 new file mode 100644 index 000000000000..6cff55e7d56f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_task_manager_optin.py @@ -0,0 +1,226 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for the gated resilient ``TaskManager`` startup recovery scan. + +``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. 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 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 + +import pytest + +from azure.ai.agentserver.core._base import ( + _has_registered_tasks, + _resilient_tasks_enabled, +) +from azure.ai.agentserver.core.tasks import ( + TaskContext, + 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 set_task_manager + + +@pytest.fixture +def _clean_state(): + """Snapshot/clear/restore the global registry, manager, and enable switch. + + ``_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_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 + finally: + _decorator_mod._REGISTERED_DESCRIPTORS.clear() + _decorator_mod._REGISTERED_DESCRIPTORS.extend(saved_desc) + set_resilient_tasks_enabled(saved_enabled) + 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 ``TaskManager`` so no real provider/token/task-store call happens.""" + _FakeTaskManager.instances.clear() + monkeypatch.setattr( + "azure.ai.agentserver.core.tasks._manager.TaskManager", + _FakeTaskManager, + ) + return _FakeTaskManager + + +def _declare_task(name: str = "gate_probe") -> None: + @task(name=name) + async def _probe(ctx: "TaskContext[dict]") -> None: + return None + + +# ------------------------------------------------------------------ # +# The two gate signals +# ------------------------------------------------------------------ # + + +class TestGateSignals: + """Unit tests for the enable switch and the descriptor-list check.""" + + def test_switch_defaults_false(self, _clean_state) -> None: + assert resilient_tasks_enabled() is False + assert _resilient_tasks_enabled() is False + + 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 + + def test_has_registered_tasks_after_task(self, _clean_state) -> None: + _declare_task() + assert _has_registered_tasks() is True + + 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 _has_registered_tasks() is True + + +# ------------------------------------------------------------------ # +# The AND gate at lifespan startup +# ------------------------------------------------------------------ # + + +class TestLifespanManagerAndRecovery: + """The TaskManager is ALWAYS constructed (cheap, no task-store calls); + 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_neither_enabled_nor_task_no_recovery( + self, _clean_state, _fake_task_manager + ) -> None: + """No switch, no task: the manager is constructed and installed so + ``get_task_manager()`` works (no ``TaskManagerNotInitialized``) — but + 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 + + app = AgentServerHost() + async with app.router.lifespan_context(app): + # 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 (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_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() + # switch left at default False + 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_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) + 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_switch_and_task_runs_recovery( + self, _clean_state, _fake_task_manager, caplog: pytest.LogCaptureFixture + ) -> None: + """Both signals true -> manager built and startup recovery runs.""" + from azure.ai.agentserver.core import AgentServerHost + + 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 + + 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 with startup recovery" in r.message for r in caplog.records) + + @pytest.mark.asyncio + 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 + + @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