dheerajturaga commented on code in PR #68039:
URL: https://github.com/apache/airflow/pull/68039#discussion_r3360020080
##########
airflow-core/src/airflow/api_fastapi/execution_api/app.py:
##########
@@ -253,9 +252,16 @@ async def _extract_w3c_trace_context(
token = otel_context.attach(ctx)
try:
yield
- finally:
- with contextlib.suppress(Exception):
- otel_context.detach(token)
+ except GeneratorExit:
+ # The dependency generator was force-closed (client disconnect /
request
+ # cancellation): the finalizer runs in a different asyncio Task — and
thus a
+ # different contextvars.Context — than attach did, so detaching the
token would
+ # raise "Token was created in a different Context" (which OTel logs at
ERROR
+ # before any suppression here could see it). The attached Context is
being
+ # discarded with the dying task, so detaching is unnecessary; skip it
and re-raise.
+ raise
+ else:
+ otel_context.detach(token)
Review Comment:
Great catch — you're right, and it was a regression. FastAPI finalizes this
dependency by `athrow`-ing the route exception into the generator at the
`yield` (confirmed via `_solve_generator` -> `asynccontextmanager` ->
`AsyncExitStack`), which took neither branch, so detach was skipped on a
same-task path where it both should and would have succeeded — leaking the
upstream W3C context into the exception handler / 500 response / error-path
spans.
Fixed by adding an `except BaseException:` arm (ordered after `except
GeneratorExit:`, since `GeneratorExit` is a `BaseException`) that detaches and
re-raises. The detach there is wrapped in `contextlib.suppress(Exception)` so a
detach failure can't mask the original route exception (preserving the
`test_route_exception_not_masked_by_detach_error` invariant). Only the
cross-task `GeneratorExit` force-close skips detach now.
---
Drafted-by: Claude Code (Opus 4.8) (no human review before posting)
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_app.py:
##########
@@ -244,3 +247,34 @@ async def mock_require_auth(request: Request) -> TIToken:
):
with pytest.raises(RuntimeError, match="boom"):
test_client.get("/variables/k", headers={"Authorization":
"Bearer fake"})
+
+ @staticmethod
+ def _make_request() -> Request:
+ return Request({"type": "http", "headers": []})
+
+ @pytest.mark.asyncio
+ async def test_detach_skipped_on_generator_exit(self):
+ """A force-closed generator (client disconnect) must NOT call detach.
+
+ Regression for the "Token was created in a different Context" error:
when the
+ dependency generator is finalised via aclose() from a different
asyncio Task,
+ detaching the token would raise and OTel would log it at ERROR. We
skip detach
+ on the GeneratorExit unwind path instead.
+ """
+ with mock.patch.object(otel_context, "detach") as detach_spy:
+ gen = _extract_w3c_trace_context(self._make_request(),
dependency_solver="fastapi")
+ await gen.asend(None) # run up to and including the yield (after
attach)
+ await gen.aclose() # raises GeneratorExit at the yield
Review Comment:
Both points addressed:
- Reworded the `test_detach_skipped_on_generator_exit` docstring to stop
claiming it reproduces the cross-context error — it closes from the same task,
so it only exercises the `GeneratorExit` branch and asserts detach is skipped;
the cross-task failure is what the spy stands in for.
- Added `test_detach_runs_on_route_exception`, which `athrow`s a
`RuntimeError` into the generator and asserts detach **is** called on that
same-task path. With the `except BaseException:` fix now detaching on route
errors, the existing `test_route_exception_not_masked_by_detach_error` is no
longer vacuous — detach is actually invoked, its mocked `ValueError` fires and
is suppressed, and the original `RuntimeError` still propagates.
---
Drafted-by: Claude Code (Opus 4.8) (no human review before posting)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
