From 81d389c90bbf151a52617bbc094474aae807f132 Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Fri, 14 Aug 2026 11:56:48 +0530 Subject: [PATCH] [agentserver] Gate blocking boot-time task recovery scan on the enablement flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The resilient TaskManager's cold-start recovery scan (Layer 1) blocks server lifespan startup on a task-store list() + credential-token acquisition. It ran whenever a durable task was registered OR the enablement switch was set. Because the responses protocol registers internal @task/@multi_turn_task primitives at host construction, every responses host — even a plain one that opted into no resilient features — paid that boot cost. Make the blocking Layer 1 scan opt-in via the existing set_resilient_tasks_enabled() switch: - TaskManager.startup() gains run_initial_scan (default True, preserving the contract for direct callers and the recovery test-suite). When False, the blocking scan is skipped; the periodic background loop (Layer 2) and request-time inline reclaim (Layer 3) still provide recovery. - The hosted lifespan passes run_initial_scan=resilient_tasks_enabled(), so a host that did not enable resilient tasks boots fast while keeping Layers 2 and 3 available. Samples: the invocations resilient samples (research, multiturn, langgraph) declare developer tasks and now call set_resilient_tasks_enabled(True) to keep eager boot recovery, matching the responses resilient samples. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../azure/ai/agentserver/core/_base.py | 22 +++++++++++++-- .../ai/agentserver/core/tasks/_manager.py | 27 ++++++++++++++++--- .../tests/tasks/test_task_manager_optin.py | 19 +++++++++---- .../samples/resilient_langgraph/app.py | 8 ++++++ .../samples/resilient_multiturn/app.py | 9 ++++++- .../samples/resilient_research/app.py | 8 ++++++ 6 files changed, 81 insertions(+), 12 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 b2f578865a8e..21ed415c8079 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 @@ -356,8 +356,26 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF set_task_manager(task_manager) if _resilient_tasks_enabled() or _has_registered_tasks(): - await task_manager.startup() - logger.info("TaskManager initialized with startup recovery") + # The blocking cold-start recovery scan (Layer 1) runs only + # when recovery was explicitly enabled. When a task subsystem + # is present but the switch is off (e.g. a plain responses + # host whose protocol registers internal primitives, or a + # developer @task app that did not force-enable), Layer 1 is + # skipped so boot is not gated on a task-store list() + + # credential-token acquisition. Recovery is still provided by + # the periodic background loop (Layer 2) and request-time + # inline reclaim (Layer 3). + _run_initial_scan = _resilient_tasks_enabled() + await task_manager.startup(run_initial_scan=_run_initial_scan) + if _run_initial_scan: + logger.info("TaskManager initialized with startup recovery") + else: + logger.info( + "TaskManager initialized (initial scan skipped; background recovery active; " + "enabled=%s, tasks_declared=%s)", + _resilient_tasks_enabled(), + _has_registered_tasks(), + ) else: logger.info( "TaskManager initialized (recovery deferred; enabled=%s, tasks_declared=%s)", 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 1cff24b64b0a..fe09a6532d41 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 @@ -725,16 +725,31 @@ async def _cancel_queued_steering_input( # pylint: disable=unused-argument if not future.done(): future.set_exception(TaskCancelled()) - async def startup(self) -> None: - """Initialize the manager and recover stale tasks. + async def startup(self, *, run_initial_scan: bool = True) -> None: + """Initialize the manager and (optionally) recover stale tasks. Called by ``AgentServerHost`` during lifespan startup. + + :keyword run_initial_scan: When True (default), run the blocking + cold-start recovery scan (Layer 1) before returning — this + preserves the historical contract relied on by direct callers + and the recovery test-suite, which assert the initial reclaim + has already happened once ``startup()`` awaits. When False, the + blocking scan is skipped so server boot is not gated on a + task-store ``list()`` + credential-token acquisition; recovery + is still provided by the periodic background loop (Layer 2, + always started here) and by request-time inline reclaim + (Layer 3). The hosted lifespan passes ``False`` unless resilient + tasks were explicitly enabled, so a plain host boots fast while + keeping Layers 2 and 3 available. + :paramtype run_initial_scan: bool """ logger.info( - "TaskManager starting (owner=%s, instance=%s, hosted=%s)", + "TaskManager starting (owner=%s, instance=%s, hosted=%s, initial_scan=%s)", self._lease_owner, self._instance_id, self._config.is_hosted, + run_initial_scan, ) # Pick up descriptors registered at import time (for recovery) from ._decorator import ( # pylint: disable=import-outside-toplevel @@ -745,7 +760,11 @@ async def startup(self) -> None: self._resume_callbacks[fn_name] = fn self._resume_opts[fn_name] = opts - await self._recover_stale_tasks() + # Layer 1: blocking cold-start recovery scan. Gated so it is + # skipped when recovery was not explicitly enabled — the periodic + # loop (Layer 2) below still covers reclaim in the background. + if run_initial_scan: + await self._recover_stale_tasks() # Layer 2: start the periodic recovery task. # Reads _PERIODIC_RECOVERY_INTERVAL_SECONDS at spawn time; 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 6cff55e7d56f..2f047b1e7154 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 @@ -69,8 +69,9 @@ def __init__(self, **kwargs) -> None: self.shutdown_called = False _FakeTaskManager.instances.append(self) - async def startup(self) -> None: + async def startup(self, *, run_initial_scan: bool = True) -> None: self.startup_called = True + self.run_initial_scan = run_initial_scan async def shutdown(self) -> None: self.shutdown_called = True @@ -163,8 +164,10 @@ async def test_neither_enabled_nor_task_no_recovery( @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).""" + """A declared task alone starts the manager's recovery machinery, but + WITHOUT the switch the blocking cold-start scan (Layer 1) is skipped — + recovery falls to the background loop (Layer 2) + inline reclaim + (Layer 3). ``startup()`` is still invoked (spawns Layer 2).""" from azure.ai.agentserver.core import AgentServerHost _declare_task() @@ -174,11 +177,14 @@ async def test_task_declared_runs_recovery_without_switch(self, _clean_state, _f pass assert len(_fake_task_manager.instances) == 1 assert _fake_task_manager.instances[0].startup_called is True + # Switch off -> Layer 1 skipped. + assert _fake_task_manager.instances[0].run_initial_scan is False @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.""" + """The switch alone runs recovery (force-enable) — the blocking + cold-start scan (Layer 1) runs and the periodic recovery loop starts + so a task declared later is picked up.""" from azure.ai.agentserver.core import AgentServerHost set_resilient_tasks_enabled(True) @@ -187,6 +193,8 @@ async def test_switch_alone_runs_recovery_without_task(self, _clean_state, _fake pass assert len(_fake_task_manager.instances) == 1 assert _fake_task_manager.instances[0].startup_called is True + # Switch on -> Layer 1 runs. + assert _fake_task_manager.instances[0].run_initial_scan is True @pytest.mark.asyncio async def test_switch_and_task_runs_recovery( @@ -206,6 +214,7 @@ async def test_switch_and_task_runs_recovery( assert len(_fake_task_manager.instances) == 1 mgr = _fake_task_manager.instances[0] assert mgr.startup_called is True + assert mgr.run_initial_scan is True assert mgr.shutdown_called is True assert any("TaskManager initialized with startup recovery" in r.message for r in caplog.records) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py index ba58ee46bdf9..3f9db93ac81a 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py @@ -67,6 +67,7 @@ EventStreamNotFoundError, streams, ) +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled from azure.ai.agentserver.invocations import InvocationAgentServerHost try: @@ -83,6 +84,13 @@ app = InvocationAgentServerHost() +# Explicitly opt into the resilient-task startup recovery scan. Declaring a +# ``@multi_turn_task`` makes the framework recovery-capable, but the eager +# boot-time reclaim of tasks orphaned by a prior crash is gated on this switch +# (background + request-time recovery still run without it). Enabling it here +# means a fresh process reclaims in-flight sessions at startup. +set_resilient_tasks_enabled(True) + async def _sse_from_stream( stream: EventStream, invocation_id: str, *, initial_status: str = "queued" diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py index b61a88ea70c9..d81db4bf3ca0 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py @@ -37,7 +37,7 @@ from starlette.responses import JSONResponse, Response from azure.ai.agentserver.core.storage import FoundryStateStore -from azure.ai.agentserver.core.tasks import TaskConflictError +from azure.ai.agentserver.core.tasks import TaskConflictError, set_resilient_tasks_enabled from azure.ai.agentserver.invocations import InvocationAgentServerHost try: @@ -47,6 +47,13 @@ app = InvocationAgentServerHost() +# Explicitly opt into the resilient-task startup recovery scan. Declaring a +# ``@multi_turn_task`` makes the framework recovery-capable, but the eager +# boot-time reclaim of tasks orphaned by a prior crash is gated on this switch +# (background + request-time recovery still run without it). Enabling it here +# means a fresh process reclaims in-flight sessions at startup. +set_resilient_tasks_enabled(True) + @app.invoke_handler async def handle_invoke(request: Request) -> Response: diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py index 00b64c2813ff..4f8b3e060f6b 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py @@ -77,6 +77,7 @@ EventStreamNotFoundError, streams, ) +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled from azure.ai.agentserver.invocations import InvocationAgentServerHost try: @@ -101,6 +102,13 @@ app = InvocationAgentServerHost() +# Explicitly opt into the resilient-task startup recovery scan. Declaring a +# ``@multi_turn_task`` makes the framework recovery-capable, but the eager +# boot-time reclaim of tasks orphaned by a prior crash is gated on this switch +# (background + request-time recovery still run without it). Enabling it here +# means a fresh process reclaims in-flight sessions at startup. +set_resilient_tasks_enabled(True) + # --- SSE rendering ---------------------------------------------------------