Show cleanup in a finally in the plain def cancellation example

Author: maxisbey

docs/handlers/cancellation.md

@@ -26,11 +26,12 @@ Put the cleanup in a `finally`:
 
 A plain `def` tool runs in a thread, and nothing can interrupt a thread from outside. The tool has to ask:
 
-```python title="server.py" hl_lines="18"
+```python title="server.py" hl_lines="22 25-26"
 --8<-- "docs_src/cancellation/tutorial002.py"
 ```
 
 * `anyio.from_thread.check_cancelled()` does nothing while the call is live, and raises once it has been cancelled. Call it between units of work.
+* Cleanup goes in a `finally` here too. Nothing in a thread awaits, so it needs no shield.
 * A `def` tool that never asks runs to the end, and its result is thrown away.
 
 ## Where it applies
@@ -48,7 +49,7 @@ It works the same over stdio and Streamable HTTP. With this SDK's `Client`, givi
 
 * When the client gives up on a call, the SDK cancels the handler: tool, prompt or resource.
 * `async def`: clean up in a `finally`, and put cleanup that awaits inside `anyio.move_on_after(seconds, shield=True)`.
-* Plain `def`: call `anyio.from_thread.check_cancelled()` between units of work, or the tool runs to the end.
+* Plain `def`: call `anyio.from_thread.check_cancelled()` between units of work, or the tool runs to the end. A plain `finally` cleans up.
 * `json_response=True` (modern connections) and `stateless_http=True` (legacy ones) switch cancellation off.
 
 Progress and cancellation are between a running tool and its *caller*. The lines it logs for *you*, the person operating the server, are a different channel: **[Logging](logging.md)**.

docs_src/cancellation/tutorial002.py

@@ -6,15 +6,21 @@
 
 mcp = MCPServer("Bookshop")
 
+offline: set[str] = set()
+
 
 def index_book(title: str) -> None:
     time.sleep(1)  # slow work with nothing to await
 
 
 @mcp.tool()
 def rebuild_index(titles: list[str]) -> str:
-    """Rebuild the search index, one book at a time."""
-    for title in titles:
-        anyio.from_thread.check_cancelled()
-        index_book(title)
-    return f"Indexed {len(titles)} books."
+    """Take search offline and rebuild its index, one book at a time."""
+    offline.add("search")
+    try:
+        for title in titles:
+            anyio.from_thread.check_cancelled()
+            index_book(title)
+        return f"Indexed {len(titles)} books."
+    finally:
+        offline.discard("search")

tests/docs_src/test_cancellation.py

@@ -20,6 +20,8 @@
 # "auto" dispatches in process; "legacy" puts a JSON-RPC stream, and so a cancellation message, in between.
 both_connections = pytest.mark.parametrize("mode", ["auto", "legacy"])
 
+TITLES = ["Dune", "Emma", "Ulysses"]
+
 
 async def abandon(
     call: Callable[[], Awaitable[object]], started: anyio.Event, then: Callable[[], object] = lambda: None
@@ -115,23 +117,35 @@ async def test_a_client_timeout_cancels_the_tool_the_same_way(
 
 @both_connections
 async def test_check_cancelled_stops_a_def_tool_at_its_next_check(mode: str, monkeypatch: pytest.MonkeyPatch) -> None:
-    """tutorial002: cancelled during the first book, the loop raises at its next check and indexes no more."""
+    """tutorial002: cancelled during the first book, the loop raises at its next check and the `finally` cleans up."""
     started = anyio.Event()
     resume = threading.Event()
     indexed: list[str] = []
 
     def index_book(title: str) -> None:
+        assert tutorial002.offline == {"search"}
         indexed.append(title)
         anyio.from_thread.run_sync(started.set)
         assert resume.wait(5)
 
     monkeypatch.setattr(tutorial002, "index_book", index_book)
-    titles = ["Dune", "Emma", "Ulysses"]
     with anyio.fail_after(5):
-        # Leaving the block waits for the tool's thread, so `indexed` is final after it.
+        # Leaving the block waits for the tool's thread, so what it left behind is final after it.
         async with Client(tutorial002.mcp, mode=mode) as client:
-            await abandon(lambda: client.call_tool("rebuild_index", {"titles": titles}), started, then=resume.set)
+            await abandon(lambda: client.call_tool("rebuild_index", {"titles": TITLES}), started, then=resume.set)
     assert indexed == ["Dune"]
+    assert tutorial002.offline == set()
+
+
+async def test_a_def_tool_nobody_cancels_indexes_every_book_and_cleans_up(monkeypatch: pytest.MonkeyPatch) -> None:
+    """tutorial002: while the call is live the checks do nothing, and the `finally` runs on a normal finish too."""
+    indexed: list[str] = []
+    monkeypatch.setattr(tutorial002, "index_book", indexed.append)
+    async with Client(tutorial002.mcp) as client:
+        result = await client.call_tool("rebuild_index", {"titles": TITLES})
+    assert result.structured_content == {"result": "Indexed 3 books."}
+    assert indexed == TITLES
+    assert tutorial002.offline == set()
 
 
 async def test_a_def_tool_that_never_checks_runs_to_the_end() -> None: