Replace plugin setup/teardown with run(server) async context manager (#3972)

CLAUDE.md

@@ -67,6 +67,7 @@ When modifying MCP functionality, changes typically need to be applied across al
 - **NEVER** create a release, comment on an issue, or open a PR unless specifically instructed to do so.
 - **NEVER** merge a PR marked as do-not-merge or draft. Check title, body, AND labels for `[DNM]`, `DNM`, `DO NOT MERGE`, `DON'T MERGE`, `DONT MERGE`, `do-not-merge`, `dont-merge`, `[DRAFT]`, or `DRAFT` (case-insensitive, any variation — some authors use `[DRAFT]` in the title even when `isDraft` is false). Authors use these as hard stops — respect them even if CI is green and review looks clean. When triaging a batch of PRs, filter these out up front AND re-check each one's labels immediately before merging, since labels can change mid-session.
 - **ALWAYS** read review-bot comments before approving a PR. CodeRabbit and chatgpt-codex-connector (Codex) leave substantive review comments on most PRs in this repo — these bots have read the diff and often flag real issues that aren't in the PR description. Use `gh pr view <num> --comments` and read the bot feedback as part of review. Unlike proposed solutions from issue reporters, review-bot feedback should be evaluated on its merits, not discounted.
+- **Be constructively skeptical of bot review comments on your own PRs.** CodeRabbit, Codex, and claude[bot] run a fresh review pass on every push, which means a PR with active churn can accumulate bot comments in a stream that never really ends — each fix surfaces a new edge case the next pass can flag. Most of the early feedback is real and worth acting on; diminishing returns set in fast. Evaluate each comment on its merits, the same way you would a human reviewer: is this a real bug users will hit, or a hypothetical that requires an adversarial setup? Does the fix introduce more complexity than the problem? Has the bot missed context that's obvious to a human reader (a `*,` keyword-only marker, a design decision documented elsewhere, something already resolved on a later commit)? When a comment is pedantic, a false positive, or flagging something already fixed, reply on the thread explaining the reasoning and move on — don't keep iterating just because more comments arrive. If you find yourself three rounds deep and the feedback is shifting toward "what if someone does X" hypotheticals, you're past the point where each fix is improving the PR. Stop, document the contract as-is, and ship.
 
 ### Releases
 

src/fastmcp/server/mixins/lifespan.py

@@ -171,14 +171,14 @@ class LifespanMixin:
             self._lifespan_result = user_lifespan_result
             self._lifespan_result_set = True
 
-            # Plugin setup pass: runs before provider lifespans and _started.
-            # Plugins may contribute providers, so this must happen before
-            # we start their lifespans below. Register teardown BEFORE the
-            # setup pass so that a partial-setup failure (one plugin's
-            # setup() raises after earlier plugins already initialized)
-            # still triggers teardown for the plugins that completed.
-            stack.push_async_callback(self._run_plugin_teardown)
-            await self._run_plugin_setup_pass()
+            # Plugin entry pass: each registered plugin's `run()` async
+            # context manager wraps the server's lifespan. Runs before
+            # provider lifespans and `_started` because plugins may
+            # contribute providers. Partial-failure safety is automatic
+            # — AsyncExitStack only unwinds plugin contexts that were
+            # successfully entered, so a raising plugin doesn't tear
+            # down plugins that never entered.
+            await self._enter_plugin_contexts(stack)
 
             # Start lifespans for all providers
             for provider in self.providers:

src/fastmcp/server/plugins/base.py

@@ -11,6 +11,8 @@ See the design document for the full specification.
 from __future__ import annotations
 
 import json
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -23,6 +25,9 @@ from fastmcp.exceptions import FastMCPError
 from fastmcp.server.middleware import Middleware
 from fastmcp.server.providers import Provider
 from fastmcp.server.transforms import Transform
+from fastmcp.utilities.logging import get_logger
+
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from starlette.routing import BaseRoute
@@ -213,23 +218,60 @@ class Plugin:
 
     # -- lifecycle ------------------------------------------------------------
 
-    async def setup(self, server: FastMCP) -> None:
-        """Called on each lifespan cycle during the server's setup pass, before the server binds.
+    @asynccontextmanager
+    async def run(self, server: FastMCP) -> AsyncIterator[None]:
+        """Async context manager wrapping the plugin's lifetime.
 
-        Receives the server it's attaching to; may call
-        `server.add_plugin()` to register additional plugins (used by
-        loader plugins). Async so that plugins can open database
-        connections, warm HTTP clients, or otherwise perform
-        `await`-able initialization. Plugins must not assume other
-        plugins are present during their own `setup()` — the full list
-        may not yet be populated.
+        The framework enters `async with plugin.run(server):` on the
+        server's lifespan stack. Everything before the `yield` runs
+        during startup (in plugin registration order); the `yield` spans
+        the server's active lifetime; everything after the `yield` runs
+        on shutdown (in reverse registration order). Cancellation on
+        shutdown unwinds the context manager automatically.
+
+        The default implementation calls `setup(server)` before the
+        `yield` and `teardown()` after it, so plugins that just need
+        one-shot init/cleanup can keep overriding just those two
+        methods. Long-running plugins (channels, integration bridges,
+        background workers) override `run()` directly to use
+        `async with` for resource management and task groups:
+
+            @asynccontextmanager
+            async def run(self, server):
+                async with httpx.AsyncClient() as client:
+                    self.client = client
+                    yield
+        """
+        await self.setup(server)
+        try:
+            yield
+        finally:
+            try:
+                await self.teardown()
+            except Exception:
+                # Exceptions during teardown are logged, not raised, so a
+                # broken plugin can't take down the server's shutdown
+                # sequence. Plugins that want different semantics should
+                # override `run()` directly.
+                logger.exception("Plugin %r raised during teardown", self.meta.name)
+
+    async def setup(self, server: FastMCP) -> None:
+        """One-shot async initialization. Called by the default `run()`
+        before the `yield`.
+
+        Override for simple init work — compile regexes, warm caches,
+        open connections, register additional plugins from a loader. For
+        anything involving long-lived resources or background tasks,
+        override `run()` directly instead and use `async with`.
         """
 
     async def teardown(self) -> None:
-        """Called on each lifespan cycle when the server shuts down, in reverse registration order.
+        """One-shot async cleanup. Called by the default `run()` after
+        the `yield`.
 
-        Async so that plugins can close connections, flush buffers, or
-        otherwise perform `await`-able cleanup.
+        Override for simple cleanup work — close connections, flush
+        buffers. For resource management that would benefit from
+        `async with`, override `run()` directly instead.
         """
 
     # -- contribution hooks ---------------------------------------------------

src/fastmcp/server/server.py

@@ -15,6 +15,7 @@ from collections.abc import (
 )
 from contextlib import (
     AbstractAsyncContextManager,
+    AsyncExitStack,
     asynccontextmanager,
 )
 from dataclasses import replace
@@ -418,28 +419,27 @@ class FastMCP(
             self.middleware.append(DereferenceRefsMiddleware())
 
         # Plugin registry: an ordered list, populated by `add_plugin()` and
-        # `plugins=[...]`. Setup and contribution collection happen during
-        # the server's lifespan startup (see `_run_plugin_setup_pass`).
+        # `plugins=[...]`. Each plugin's `run()` async context manager wraps
+        # the server's lifespan (see `_enter_plugin_contexts`).
         self.plugins: list[Plugin] = []
         # Plugins whose contributions (middleware/transforms/providers/routes)
         # have been collected onto the server. Server contributions persist
         # across lifespan cycles like server-authored middleware, so we
         # don't re-install them on re-entry.
         self._plugins_contributed: set[int] = set()
-        # Plugins whose setup() has completed successfully in the current
-        # lifespan cycle. Reset on teardown. Used to scope teardown to
-        # plugins that actually ran setup, so a partial-setup failure
-        # still triggers teardown for everything that was initialized.
-        self._plugins_set_up: set[int] = set()
+        # Plugins whose `run()` context has been entered in the current
+        # lifespan cycle. Used to dedupe entry when the same instance is
+        # registered twice; reset after the cycle's ephemeral cleanup.
+        self._plugins_entered: set[int] = set()
         # Per-plugin record of contributions we installed, stored as
         # (container, item) tuples so we can reverse them when an
         # ephemeral plugin is torn down.
         self._plugin_contributions: dict[int, list[tuple[list[Any], Any]]] = {}
-        # True while the setup pass is executing. `add_plugin()` uses this
-        # to flag ephemeral plugins — plugins added from inside another
-        # plugin's setup() (the loader pattern). Ephemeral plugins are
-        # removed on teardown so loaders can freshly re-hydrate children
-        # on the next lifespan cycle without accumulating duplicates.
+        # True while the plugin-entry loop is running. `add_plugin()` uses
+        # this to flag ephemeral plugins — plugins added from inside
+        # another plugin's `run()` (the loader pattern). Ephemeral plugins
+        # are removed after all plugin contexts exit so loaders can
+        # freshly re-hydrate children on the next lifespan cycle.
         self._in_plugin_setup_pass: bool = False
         for p in plugins or []:
             self.add_plugin(p)
@@ -562,110 +562,102 @@ class FastMCP(
             self._additional_http_routes.append(route)
             records.append((self._additional_http_routes, route))
 
-    async def _run_plugin_setup_pass(self) -> None:
-        """Run setup() on every registered plugin and collect contributions.
+    async def _enter_plugin_contexts(self, stack: AsyncExitStack) -> None:
+        """Enter each registered plugin's `run()` context on the given stack.
 
         Called during server startup (from `_lifespan_manager`), before
-        the server binds. Iterates the plugin list in order, awaiting
-        `setup(server)` on each; plugins added during another plugin's
-        setup (the loader pattern) are picked up by the same loop because
-        the iteration advances against a live index.
+        the server binds. Iterates the plugin list in order, entering
+        `async with plugin.run(server):` on the shared exit stack for
+        each plugin. Plugins added during another plugin's `run()` (the
+        loader pattern) are picked up by the same loop because the
+        iteration advances against a live index.
 
-        Setup runs every lifespan cycle — plugins expect a matching
-        `setup`/`teardown` pair. Contribution collection is one-shot
-        per plugin: once a plugin's middleware/transforms/providers/routes
-        have been installed on the server they persist across cycles,
-        matching how server-authored middleware behaves.
+        Contributions (middleware, transforms, providers) are collected
+        once per plugin across all lifespan cycles — once installed, they
+        persist on the server just like server-authored middleware.
+        Routes were already collected synchronously at `add_plugin()`
+        time so HTTP transports see them before the lifespan runs.
+
+        Ephemeral cleanup is registered as a post-stack callback so it
+        runs after every plugin's `run()` has exited but before the
+        outer server lifespan unwinds further.
         """
-        # Setup pass: mutating-list iteration. New plugins appended by a
-        # plugin's setup() are picked up on subsequent iterations. We mark
-        # each plugin as "set up" only after setup() returns so that a
-        # partial-setup failure scopes teardown to plugins that actually
-        # completed initialization. The _in_plugin_setup_pass flag lets
-        # add_plugin() mark new plugins as ephemeral (see add_plugin).
+        # Register ephemeral cleanup FIRST so it unwinds LAST (after all
+        # plugin.run() contexts have exited).
+        stack.push_async_callback(self._cleanup_ephemeral_plugins)
+
+        # Plugin-entry loop: mutating-list iteration. New plugins appended
+        # by a plugin's run() (loader pattern) are picked up on subsequent
+        # iterations. `_in_plugin_setup_pass` lets add_plugin() mark those
+        # as ephemeral (see add_plugin). `_plugins_entered` dedupes: a
+        # plugin instance registered twice only enters its run() once per
+        # lifespan cycle, keeping setup/teardown counts symmetric.
         self._in_plugin_setup_pass = True
         try:
             i = 0
             while i < len(self.plugins):
                 plugin = self.plugins[i]
-                await plugin.setup(self)
-                self._plugins_set_up.add(id(plugin))
                 i += 1
-
-            # Contribution collection: run in registration order. Guarded
-            # per-plugin because contributions persist across lifespan
-            # cycles; new plugins (e.g. added by a loader during this
-            # cycle's setup) still get their contributions collected. Note:
-            # routes are collected synchronously at add_plugin() time
-            # because HTTP transports snapshot the route list before the
-            # lifespan runs.
-            for plugin in self.plugins:
-                if id(plugin) in self._plugins_contributed:
+                if id(plugin) in self._plugins_entered:
                     continue
-                # Gather everything first: any hook that raises aborts
-                # before any server state is mutated, so a retry on the
-                # next lifespan cycle starts clean rather than appending
-                # duplicate middleware on top of partial contributions.
-                mws = list(plugin.middleware())
-                transforms_ = list(plugin.transforms())
-                providers_ = list(plugin.providers())
-
-                records = self._plugin_contributions.setdefault(id(plugin), [])
-                for mw in mws:
-                    self.add_middleware(mw)
-                    records.append((self.middleware, mw))
-                for transform in transforms_:
-                    self.add_transform(transform)
-                    records.append((self._transforms, transform))
-                for provider in providers_:
-                    # add_provider may wrap the value (for example a
-                    # FastMCP is wrapped in FastMCPProvider). Record
-                    # whatever actually landed in self.providers so
-                    # teardown can find it by identity.
-                    before = len(self.providers)
-                    self.add_provider(provider)
-                    for stored in self.providers[before:]:
-                        records.append((self.providers, stored))
-                self._plugins_contributed.add(id(plugin))
+                self._plugins_entered.add(id(plugin))
+                try:
+                    await stack.enter_async_context(plugin.run(self))
+                except Exception:
+                    logger.exception(
+                        "Plugin %r raised while entering run()", plugin.meta.name
+                    )
+                    raise
         finally:
             self._in_plugin_setup_pass = False
 
-    async def _run_plugin_teardown(self) -> None:
-        """Await `teardown()` on plugins whose setup() completed this cycle.
-
-        Iterates the plugin list in reverse registration order and only
-        tears down plugins that were successfully set up. This makes
-        teardown safe to register before `_run_plugin_setup_pass` — a
-        failure partway through setup still runs teardown for the plugins
-        that had already initialized.
-
-        Ephemeral plugins (those added from inside another plugin's
-        setup() by a loader) are removed from the server after teardown
-        runs, along with the contributions they installed. This keeps
-        loader plugins from accumulating duplicate children across
-        repeated lifespan cycles.
-        """
-        # Snapshot + clear first so that a slow teardown combined with a
-        # re-entry cannot double-count the set.
-        set_up = self._plugins_set_up
-        self._plugins_set_up = set()
-        for plugin in reversed(self.plugins):
-            if id(plugin) not in set_up:
+        # Contribution collection: run in registration order. Guarded
+        # per-plugin because contributions persist across lifespan cycles;
+        # new plugins (e.g. added by a loader during this cycle's run)
+        # still get their contributions collected. Routes are collected
+        # synchronously at add_plugin() time (see above).
+        for plugin in self.plugins:
+            if id(plugin) in self._plugins_contributed:
                 continue
-            # Discard first so a plugin that appears twice in the list
-            # (same instance registered twice — see test_duplicates_allowed)
-            # only gets torn down once, even though its setup() ran once
-            # per list entry.
-            set_up.discard(id(plugin))
-            try:
-                await plugin.teardown()
-            except Exception:
-                logger.exception("Plugin %r raised during teardown", plugin.meta.name)
+            # Gather everything first: any hook that raises aborts before
+            # any server state is mutated, so a retry on the next lifespan
+            # cycle starts clean rather than appending duplicate middleware
+            # on top of partial contributions.
+            mws = list(plugin.middleware())
+            transforms_ = list(plugin.transforms())
+            providers_ = list(plugin.providers())
+
+            records = self._plugin_contributions.setdefault(id(plugin), [])
+            for mw in mws:
+                self.add_middleware(mw)
+                records.append((self.middleware, mw))
+            for transform in transforms_:
+                self.add_transform(transform)
+                records.append((self._transforms, transform))
+            for provider in providers_:
+                # add_provider may wrap the value (for example a FastMCP
+                # is wrapped in FastMCPProvider). Record whatever
+                # actually landed in self.providers so teardown can find
+                # it by identity.
+                before = len(self.providers)
+                self.add_provider(provider)
+                for stored in self.providers[before:]:
+                    records.append((self.providers, stored))
+            self._plugins_contributed.add(id(plugin))
+
+    async def _cleanup_ephemeral_plugins(self) -> None:
+        """Remove ephemeral (loader-added) plugins after all contexts have exited.
+
+        Runs as an `AsyncExitStack` callback registered in
+        `_enter_plugin_contexts` so it fires after every plugin's `run()`
+        has unwound. Without this, each lifespan cycle would accumulate a
+        fresh generation of loader-added children in `self.plugins`, and
+        their contributions would accumulate in the server's
+        middleware/transform/provider lists.
+        """
+        # Reset the per-cycle entered set for the next cycle.
+        self._plugins_entered = set()
 
-        # Remove ephemeral plugins and their contributions. On the next
-        # lifespan cycle, the loader that produced them will freshly
-        # re-hydrate its children — without this cleanup, the plugin list
-        # and contribution registries would grow on every cycle.
         ephemeral = [p for p in self.plugins if getattr(p, "_fastmcp_ephemeral", False)]
         for plugin in ephemeral:
             records = self._plugin_contributions.pop(id(plugin), [])

tests/server/test_plugins.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
+from contextlib import suppress
 from pathlib import Path
 
 import pytest
@@ -692,6 +694,159 @@ class TestLifecycle:
         assert mcp.middleware == baseline_middleware
 
 
+class TestRunHook:
+    """Plugins that override `run()` directly (the long-running pattern)."""
+
+    async def test_run_override_wraps_server_lifetime(self):
+        """A plugin overriding run() sees the server live between setup and teardown."""
+        from contextlib import asynccontextmanager
+
+        recorder = _Recorder()
+
+        class Long(Plugin):
+            meta = PluginMeta(name="long", version="0.1.0")
+
+            @asynccontextmanager
+            async def run(self, server):
+                recorder.events.append(("enter", "long"))
+                try:
+                    yield
+                finally:
+                    recorder.events.append(("exit", "long"))
+
+        mcp = FastMCP("t", plugins=[Long()])
+        async with Client(mcp) as c:
+            await c.ping()
+            # Mid-cycle: enter fired, exit hasn't.
+            assert ("enter", "long") in recorder.events
+            assert ("exit", "long") not in recorder.events
+
+        # After teardown: both fired.
+        assert recorder.events == [("enter", "long"), ("exit", "long")]
+
+    async def test_run_override_can_use_async_with(self):
+        """A plugin's run() can acquire an async-context resource and release it on exit."""
+        from contextlib import asynccontextmanager
+
+        recorder = _Recorder()
+
+        @asynccontextmanager
+        async def fake_resource():
+            recorder.events.append(("acquire", "resource"))
+            try:
+                yield "handle"
+            finally:
+                recorder.events.append(("release", "resource"))
+
+        class WithResource(Plugin):
+            meta = PluginMeta(name="with-resource", version="0.1.0")
+
+            @asynccontextmanager
+            async def run(self, server):
+                async with fake_resource() as handle:
+                    self.handle = handle
+                    yield
+
+        p = WithResource()
+        mcp = FastMCP("t", plugins=[p])
+        async with Client(mcp) as c:
+            await c.ping()
+            assert p.handle == "handle"
+
+        # async with cleanup fired on exit path
+        assert recorder.events == [
+            ("acquire", "resource"),
+            ("release", "resource"),
+        ]
+
+    async def test_run_override_cancellation_propagates_into_background_task(self):
+        """A long-running background task inside run() is cancelled on shutdown."""
+        from contextlib import asynccontextmanager
+
+        recorder = _Recorder()
+
+        class Background(Plugin):
+            meta = PluginMeta(name="background", version="0.1.0")
+
+            @asynccontextmanager
+            async def run(self, server):
+                async def worker():
+                    try:
+                        await asyncio.Event().wait()
+                    except asyncio.CancelledError:
+                        recorder.events.append(("cancelled", "worker"))
+                        raise
+
+                task = asyncio.create_task(worker())
+                try:
+                    yield
+                finally:
+                    task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await task
+
+        mcp = FastMCP("t", plugins=[Background()])
+        async with Client(mcp) as c:
+            await c.ping()
+
+        assert recorder.events == [("cancelled", "worker")]
+
+    async def test_run_override_raising_before_yield_aborts_startup(self):
+        """If a plugin's run() raises before yielding, startup fails cleanly."""
+        from contextlib import asynccontextmanager
+
+        class BadStart(Plugin):
+            meta = PluginMeta(name="bad-start", version="0.1.0")
+
+            @asynccontextmanager
+            async def run(self, server):
+                raise RuntimeError("cannot start")
+                yield  # unreachable
+
+        mcp = FastMCP("t", plugins=[BadStart()])
+        with pytest.raises(RuntimeError, match="cannot start"):
+            async with Client(mcp) as c:
+                await c.ping()
+
+    async def test_run_override_composes_with_simple_setup_teardown_plugins(self):
+        """A server can mix run-override plugins with setup/teardown plugins."""
+        from contextlib import asynccontextmanager
+
+        recorder = _Recorder()
+
+        class Simple(Plugin):
+            meta = PluginMeta(name="simple", version="0.1.0")
+
+            async def setup(self, server):
+                recorder.events.append(("setup", "simple"))
+
+            async def teardown(self):
+                recorder.events.append(("teardown", "simple"))
+
+        class LongRunning(Plugin):
+            meta = PluginMeta(name="long-running", version="0.1.0")
+
+            @asynccontextmanager
+            async def run(self, server):
+                recorder.events.append(("enter", "long-running"))
+                try:
+                    yield
+                finally:
+                    recorder.events.append(("exit", "long-running"))
+
+        mcp = FastMCP("t", plugins=[Simple(), LongRunning()])
+        async with Client(mcp) as c:
+            await c.ping()
+
+        # Enter order follows registration; exit order is reversed.
+        assert recorder.events == [
+            ("setup", "simple"),
+            ("enter", "long-running"),
+            ("exit", "long-running"),
+            ("teardown", "simple"),
+        ]
+
+
 class TestContributions:
     """Plugin contributions are installed during the setup pass."""