Nathandrake229 · Oct 20, 2026 · Oct 20, 2026 · Oct 20, 2026
@@ -4,6 +4,11 @@
 
 ### Features Added
 
+- Added protocol-scoped task-runtime startup requests so protocol hosts can
+  obtain scoped `TaskManager` access without mutating the process-wide
+  resilient-task switch. This prevents one protocol's enablement from
+  activating another protocol in a mixed host.
+
 ### Breaking Changes
 
 ### Bugs Fixed

@@ -25,7 +25,7 @@
 from . import _config, _tracing
 from ._errors import create_error_response
 from ._lifecycle import AgentSessionContext, _LifecycleState
-from ._middleware import InboundRequestLoggingMiddleware
+from ._middleware import InboundRequestLoggingMiddleware, _TaskManagerMiddleware
 from ._request_context import (
     FoundryAgentRequestContext,
     get_request_context,
@@ -273,6 +273,11 @@ def __init__(  # pylint: disable=too-many-statements
         self._shutdown_fn: Optional[Callable[[], Awaitable[None]]] = None
         self._readiness_fn: Optional[Callable[[Request], Awaitable[Response]]] = None
         self._lifecycle = _LifecycleState()
+        # Protocol mixins contribute independent startup requirements via cooperative init.
+        task_manager_requirements: list[Callable[[], bool]] = kwargs.pop(
+            "_task_manager_requirements", [_resilient_tasks_enabled]
+        )
+        self._task_manager: Any = None
         #  Pre-shutdown callbacks invoked SYNCHRONOUSLY from the
         # SIGTERM signal handler — before Hypercorn's graceful drain
         # begins. Used by responses to set ``_shutdown_requested`` early so
@@ -316,7 +321,7 @@ def __init__(  # pylint: disable=too-many-statements
 
         # Build lifespan context manager
         @contextlib.asynccontextmanager
-        async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF029
+        async def _run_lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF029
             logger.info("AgentServerHost started")
 
             # --- Startup configuration logging ---
@@ -345,26 +350,10 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF
                 protocols,
             )
 
-            # --- Resilient task manager initialization ---
-            #
-            # The resilient ``TaskManager`` is CONSTRUCTED ONLY when resilient
-            # tasks were explicitly enabled via
-            # ``set_resilient_tasks_enabled(True)``. Recovery (and the durable
-            # task subsystem as a whole) is strictly opt-in: a host that does
-            # not set the switch pays nothing — no manager, no provider client,
-            # no task-store call, no recovery loop.
-            #
-            # When the switch is OFF the manager is NOT installed, so
-            # ``get_task_manager()`` raises ``TaskManagerNotInitialized``.
-            # Callers that route through the manager (e.g. the responses
-            # protocol's ``store=true`` path) are expected to SWALLOW that
-            # signal and degrade to non-durable, in-process execution — the
-            # response still runs and persists, it simply is not crash-
-            # recoverable. Declaring a durable task (``@task`` /
-            # ``@multi_turn_task``) does NOT implicitly turn the subsystem on;
-            # an app that wants durable tasks / recovery must set the switch.
+            # Invocations uses the core switch; Responses contributes its own
+            # feature flags. Merely registering task callbacks is not an opt-in.
             task_manager: Any = None
-            if _resilient_tasks_enabled():
+            if any(required() for required in task_manager_requirements):
                 # Resilient tasks were explicitly enabled. Construct the manager
                 # and run startup recovery. If EITHER fails, durability was
                 # explicitly requested, so fail the lifespan (fail-fast at boot)
@@ -381,6 +370,7 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF
                     shutdown_event=asyncio.Event(),
                     shutdown_grace_seconds=_read_task_manager_shutdown_grace(),
                 )
+                self._task_manager = task_manager
                 set_task_manager(task_manager)
                 try:
                     await task_manager.startup()
@@ -395,14 +385,17 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF
                         await task_manager.shutdown()
                     except Exception:  # pylint: disable=broad-exception-caught
                         logger.warning("Error shutting down TaskManager after startup failure", exc_info=True)
-                    set_task_manager(None)
+                    from .tasks._manager import _clear_task_manager  # pylint: disable=import-outside-toplevel
+
+                    _clear_task_manager(task_manager)
+                    self._task_manager = None
                     task_manager = None
                     raise
                 logger.info("TaskManager initialized with startup recovery")
             else:
                 logger.info(
-                    "TaskManager NOT initialized (resilient tasks disabled; enable via "
-                    "set_resilient_tasks_enabled(True)). Durable tasks and crash recovery "
+                    "TaskManager NOT initialized (no protocol requested task execution). "
+                    "Durable tasks and crash recovery "
                     "are inactive; store=true work degrades to non-durable execution. "
                     "tasks_declared=%s",
                     _has_registered_tasks(),
@@ -445,15 +438,24 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF
             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)
 
+        @contextlib.asynccontextmanager
+        async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
+            from .tasks._manager import _task_manager_scope  # pylint: disable=import-outside-toplevel
+
+            with _task_manager_scope(lambda: self._task_manager):
+                try:
+                    async with _run_lifespan(app):
+                        yield
+                finally:
+                    from .tasks._manager import _clear_task_manager  # pylint: disable=import-outside-toplevel
+
+                    _clear_task_manager(self._task_manager)
+                    self._task_manager = None
+
         async def _before_snapshot_route(request: Request) -> Response:
             return await AgentServerHost._before_snapshot_endpoint(self, request)
 
@@ -498,6 +500,7 @@ async def _after_restore_route(request: Request) -> Response:
             routes=all_routes,
             lifespan=_lifespan,
             middleware=[
+                Middleware(_TaskManagerMiddleware, get_manager=lambda: self._task_manager),
                 Middleware(InboundRequestLoggingMiddleware),  # type: ignore[arg-type]
                 Middleware(  # type: ignore[arg-type]
                     _PlatformHeaderMiddleware,
@@ -737,14 +740,10 @@ def _parse_session_context(payload: dict[str, Any]) -> AgentSessionContext:
         if not isinstance(restore_id, str) or not restore_id:
             raise ValueError("restore_id must be a non-empty string")
         if not isinstance(raw_overrides, dict) or any(
-            not isinstance(key, str) or not isinstance(value, str)
-            for key, value in raw_overrides.items()
+            not isinstance(key, str) or not isinstance(value, str) for key, value in raw_overrides.items()
         ):
             raise ValueError("session_env_overrides must map strings to strings")
-        if any(
-            not key or "=" in key or "\0" in key or "\0" in value
-            for key, value in raw_overrides.items()
-        ):
+        if any(not key or "=" in key or "\0" in key or "\0" in value for key, value in raw_overrides.items()):
             raise ValueError("session_env_overrides contains an invalid environment entry")
         override_session_id = raw_overrides.get(_SESSION_ID_ENV)
         if override_session_id is not None and override_session_id != session_id:
@@ -780,8 +779,7 @@ async def _before_snapshot_endpoint(self, request: Request) -> Response:
                     status_code=500,
                 )
             self._lifecycle.captured_environment_values = {
-                key: os.environ.get(key)
-                for key in self._lifecycle.applied_environment_variables
+                key: os.environ.get(key) for key in self._lifecycle.applied_environment_variables
             }
             self._lifecycle.before_snapshot_completed = True
             return JSONResponse({"status": "ok"})
@@ -815,20 +813,12 @@ async def _after_restore_endpoint(self, request: Request) -> Response:
 
             effective_environment = dict(context.session_env_overrides)
             effective_environment.setdefault(_SESSION_ID_ENV, context.session_id)
-            environment_variables = (
-                self._lifecycle.applied_environment_variables
-                | effective_environment.keys()
-            )
-            previous_environment = {
-                key: os.environ.get(key, _MISSING_ENV)
-                for key in environment_variables
-            }
+            environment_variables = self._lifecycle.applied_environment_variables | effective_environment.keys()
+            previous_environment = {key: os.environ.get(key, _MISSING_ENV) for key in environment_variables}
             captured_environment = dict(self._lifecycle.captured_environment_values)
             for key, value in previous_environment.items():
                 if key not in captured_environment:
-                    captured_environment[key] = (
-                        None if value is _MISSING_ENV else str(value)
-                    )
+                    captured_environment[key] = None if value is _MISSING_ENV else str(value)
             previous_config = vars(self.config).copy()
             task_manager = _get_active_task_manager()
 
@@ -842,9 +832,7 @@ async def _after_restore_endpoint(self, request: Request) -> Response:
                 restored_config = _config.AgentConfig.from_env()
                 vars(self.config).update(vars(restored_config))
                 if task_manager is not None:
-                    task_manager._rehydrate_session_state(  # pylint: disable=protected-access
-                        context.restore_id
-                    )
+                    task_manager._rehydrate_session_state(context.restore_id)  # pylint: disable=protected-access
 
                 current_request_context = get_request_context()
                 request_context_token = set_request_context(

@@ -20,13 +20,51 @@
 
 import logging
 import time
-from typing import Any, MutableMapping
+from typing import Any, Callable, MutableMapping
 
 from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.middleware import Middleware
 
 logger = logging.getLogger("azure.ai.agentserver")
 
 
+class _TaskManagerMiddleware:
+    """Scope task access to the host and protocol serving this request."""
+
+    def __init__(self, app: ASGIApp, *, get_manager: Callable[[], Any]) -> None:
+        self.app = app
+        self._get_manager = get_manager
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        from .tasks._manager import _task_manager_scope  # pylint: disable=import-outside-toplevel
+
+        with _task_manager_scope(self._get_manager):
+            await self.app(scope, receive, send)
+
+
+def _task_runtime_middleware(host: Any, kwargs: dict[str, Any], enabled: Callable[[], bool]) -> list[Middleware]:
+    """Register protocol activation and scope its routes to the host's runtime.
+
+    :param host: The protocol host being constructed.
+    :type host: Any
+    :param kwargs: Cooperative host-constructor arguments.
+    :type kwargs: dict[str, Any]
+    :param enabled: Whether this protocol requests task execution.
+    :type enabled: Callable[[], bool]
+    :return: Middleware for this protocol's routes.
+    :rtype: list[~starlette.middleware.Middleware]
+    """
+    requirements = list(kwargs.pop("_task_manager_requirements", []))
+    requirements.append(enabled)
+    kwargs["_task_manager_requirements"] = requirements
+    return [
+        Middleware(
+            _TaskManagerMiddleware,
+            get_manager=lambda: host._task_manager if enabled() else None,  # pylint: disable=protected-access
+        )
+    ]
+
+
 def _extract_header(headers: list[tuple[bytes, bytes]], name: bytes) -> str | None:
     """Extract a header value from raw ASGI headers.
 
@@ -147,7 +185,10 @@ async def _send_wrapper(message: MutableMapping[str, Any]) -> None:
             elapsed_ms = (time.monotonic() - start) * 1000
             logger.warning(
                 "Inbound %s %s failed with status 500 in %.1fms%s",
-                method, path, elapsed_ms, extra_str,
+                method,
+                path,
+                elapsed_ms,
+                extra_str,
             )
             raise
 
@@ -156,10 +197,18 @@ async def _send_wrapper(message: MutableMapping[str, Any]) -> None:
         if status_code is not None and status_code >= 400:
             logger.warning(
                 "Inbound %s %s completed with status %d in %.1fms%s",
-                method, path, status_code, elapsed_ms, extra_str,
+                method,
+                path,
+                status_code,
+                elapsed_ms,
+                extra_str,
             )
         else:
             logger.info(
                 "Inbound %s %s completed with status %s in %.1fms%s",
-                method, path, status_code, elapsed_ms, extra_str,
+                method,
+                path,
+                status_code,
+                elapsed_ms,
+                extra_str,
             )
@@ -36,6 +36,7 @@ async def my_task(ctx: TaskContext[MyInput]) -> MyOutput:
 
 from ._client import TransportClassifiedError as _TransportClassifiedError
 from ._context import TaskContext
+from ._exceptions import TaskManagerNotInitialized
 from ._exceptions_internal import _HostedConflict, _translate_hosted_conflict
 from ._retry import RetryPolicy
 from ._run import TaskRun
@@ -478,8 +479,10 @@ def __init__(
         # resume tables so _recover_stale_tasks / get_active_run can pick
         # up the multi-turn opts (_is_multi_turn).
         try:
-            from ._manager import _manager as _live_manager  # pylint: disable=import-outside-toplevel
-        except ImportError:  # pragma: no cover
+            from ._manager import get_task_manager  # pylint: disable=import-outside-toplevel
+
+            _live_manager = get_task_manager()
+        except (ImportError, TaskManagerNotInitialized):
             _live_manager = None  # type: ignore[assignment]
         if _live_manager is not None:
             try:
@@ -1798,28 +1801,28 @@ def multi_turn_task(
 ) -> Any:
     """Decorator producing a multi-turn resilient chain.
 
-        Multi-turn chains accept inputs across many turns against the same
-        ``task_id``. The handler's ``return X`` is the implicit-suspend
-        signal — there is no ``ctx.suspend``. The chain stays
-        alive across handler raises.
-
-        :param fn: The handler to decorate (parameterized form only; the bare
-            ``@multi_turn_task`` form is rejected because it cannot supply a
-            ``name``).
-        :type fn: Callable[..., Any] | None
-        :keyword name: **Required** stable chain-identity anchor (no
-            function-derived default; omitting it raises ``ValueError``).
-        :keyword title: Static human-readable string. Callable-factory form is
-            not supported.
-        :keyword timeout: Per-turn cooperative timeout. Defaults to 1 day when
-            unset; 7 days is a hard ceiling (larger/negative rejected at
-            decoration). This is a per-turn cap, not a chain-lifetime bound.
-        :keyword retry: Default retry policy.
-        :keyword steerable: When True, ``start()`` against an in-flight chain
-            queues the new input instead of raising ``TaskConflictError``.
-        :return: A :class:`MultiTurnTask` instance (distinct public class from
-            :class:`Task`).
-        :rtype: Any
+    Multi-turn chains accept inputs across many turns against the same
+    ``task_id``. The handler's ``return X`` is the implicit-suspend
+    signal — there is no ``ctx.suspend``. The chain stays
+    alive across handler raises.
+
+    :param fn: The handler to decorate (parameterized form only; the bare
+        ``@multi_turn_task`` form is rejected because it cannot supply a
+        ``name``).
+    :type fn: Callable[..., Any] | None
+    :keyword name: **Required** stable chain-identity anchor (no
+        function-derived default; omitting it raises ``ValueError``).
+    :keyword title: Static human-readable string. Callable-factory form is
+        not supported.
+    :keyword timeout: Per-turn cooperative timeout. Defaults to 1 day when
+        unset; 7 days is a hard ceiling (larger/negative rejected at
+        decoration). This is a per-turn cap, not a chain-lifetime bound.
+    :keyword retry: Default retry policy.
+    :keyword steerable: When True, ``start()`` against an in-flight chain
+        queues the new input instead of raising ``TaskConflictError``.
+    :return: A :class:`MultiTurnTask` instance (distinct public class from
+        :class:`Task`).
+    :rtype: Any
     """
     #  — reject unknown kwargs at decoration time
     _validate_multi_turn_task_kwargs(**_extra_kwargs)