From d4cb8b464bf1cbe69a8921fc8c9315e04a5f49cb Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 12 Jun 2026 13:21:58 -0700
Subject: [PATCH] Fix for trace_child_tasks exception handling

I had Claude Fable 5 review our use of contextvar and
it spotted this place where exceptions were not
correctly handled.
---
 datasette/tracer.py  |  6 ++++--
 tests/test_tracer.py | 13 +++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/datasette/tracer.py b/datasette/tracer.py
index 9e66613b..28f3cc09 100644
--- a/datasette/tracer.py
+++ b/datasette/tracer.py
@@ -27,8 +27,10 @@ def get_task_id():
 @contextmanager
 def trace_child_tasks():
     token = trace_task_id.set(get_task_id())
-    yield
-    trace_task_id.reset(token)
+    try:
+        yield
+    finally:
+        trace_task_id.reset(token)
 
 
 @contextmanager
diff --git a/tests/test_tracer.py b/tests/test_tracer.py
index 6cc80fc4..9db211d3 100644
--- a/tests/test_tracer.py
+++ b/tests/test_tracer.py
@@ -70,6 +70,19 @@ def test_trace_query_errors():
     assert trace_info["traces"][-1]["error"] == "no such table: non_existent_table"
 
 
+@pytest.mark.asyncio
+async def test_trace_child_tasks_resets_contextvar_on_exception():
+    from datasette import tracer
+
+    before = tracer.trace_task_id.get()
+    with pytest.raises(ValueError):
+        with tracer.trace_child_tasks():
+            assert tracer.trace_task_id.get() is not None
+            raise ValueError("simulated error")
+    # The contextvar must be reset even though the block raised
+    assert tracer.trace_task_id.get() == before
+
+
 def test_trace_parallel_queries():
     with make_app_client(settings={"trace_debug": True}) as client:
         response = client.get("/parallel-queries?_trace=1")