Repository: agronholm/anyio

Issue:
Assertion error when `current_task()` returns an instance of `_asyncio.Task`
### Things to check first

- [X] I have searched the existing issues and didn't find my bug already reported there

- [X] I have checked that my bug is still present in the latest release


### AnyIO version

4.7.0

### Python version

3.12.8

### What happened?

It seems that sometimes when `asyncio.current_task()` is called in

https://github.com/agronholm/anyio/blob/97f064295b65dc1d7a2ded6fd24ba5be46d4c4f4/src/anyio/_backends/_asyncio.py#L2367

an instance of `_asyncio.Task` is returned. This then causes an assertion to be raised in

https://github.com/agronholm/anyio/blob/97f064295b65dc1d7a2ded6fd24ba5be46d4c4f4/src/anyio/_backends/_asyncio.py#L685-L694


### How can we reproduce the bug?

This is triggered in our test suite (fastapi, httpx, pytest-asyncio) after upgrading to anyio 4.7.0, but I haven't been able to come up with a MRE so hopefully this is enough to point maintainers to the potential fix.

Discussion:
agronholm: I don't know how to write a test for this, as `nest_asyncio.apply()` is too invasive to be used in the test suite.
graingert: you could try with:

```
with mock.patch("asyncio.Task", asyncio.tasks._PyTask):
    ...
```
agronholm: I've added a test and verified that it fails without this patch.
edgarrmondragon: I can confirm that this patch fixes the problem for us
sigridjineth: @agronholm I just have came across this, but would be better implementing the logic for accomodating some varities of task-like objects I think. what's your opinion?
agronholm: > @agronholm I just have came across this, but would be better implementing the logic for accomodating some varities of task-like objects I think. what's your opinion?

If you find a problem not fixed by this PR, let me know. Otherwise I'm not inclined to add any more code just because.

Patch:
diff --git a/docs/versionhistory.rst b/docs/versionhistory.rst
index 518aef8..4901da4 100644
--- a/docs/versionhistory.rst
+++ b/docs/versionhistory.rst
@@ -9,6 +9,8 @@ This library adheres to `Semantic Versioning 2.0 <http://semver.org/>`_.
   thread waker socket pair. This should improve the performance of ``wait_readable()``
   and ``wait_writable()`` when using the ``ProactorEventLoop``
   (`#836 <https://github.com/agronholm/anyio/pull/836>`_; PR by @graingert)
+- Fixed ``AssertionError`` when using ``nest-asyncio``
+  (`#840 <https://github.com/agronholm/anyio/issues/840>`_)
 
 **4.7.0**
 
diff --git a/src/anyio/_backends/_asyncio.py b/src/anyio/_backends/_asyncio.py
index 0b7479d..5a0aa93 100644
--- a/src/anyio/_backends/_asyncio.py
+++ b/src/anyio/_backends/_asyncio.py
@@ -677,40 +677,42 @@ class TaskState:
         self.cancel_scope = cancel_scope
 
 
-class TaskStateStore(MutableMapping["Awaitable[Any] | asyncio.Task", TaskState]):
+class TaskStateStore(
+    MutableMapping["Coroutine[Any, Any, Any] | asyncio.Task", TaskState]
+):
     def __init__(self) -> None:
         self._task_states = WeakKeyDictionary[asyncio.Task, TaskState]()
-        self._preliminary_task_states: dict[Awaitable[Any], TaskState] = {}
+        self._preliminary_task_states: dict[Coroutine[Any, Any, Any], TaskState] = {}
 
-    def __getitem__(self, key: Awaitable[Any] | asyncio.Task, /) -> TaskState:
-        assert isinstance(key, asyncio.Task)
+    def __getitem__(self, key: Coroutine[Any, Any, Any] | asyncio.Task, /) -> TaskState:
+        task = cast(asyncio.Task, key)
         try:
-            return self._task_states[key]
+            return self._task_states[task]
         except KeyError:
-            if coro := key.get_coro():
+            if coro := task.get_coro():
                 if state := self._preliminary_task_states.get(coro):
                     return state
 
         raise KeyError(key)
 
     def __setitem__(
-        self, key: asyncio.Task | Awaitable[Any], value: TaskState, /
+        self, key: asyncio.Task | Coroutine[Any, Any, Any], value: TaskState, /
     ) -> None:
-        if isinstance(key, asyncio.Task):
-            self._task_states[key] = value
-        else:
+        if isinstance(key, Coroutine):
             self._preliminary_task_states[key] = value
-
-    def __delitem__(self, key: asyncio.Task | Awaitable[Any], /) -> None:
-        if isinstance(key, asyncio.Task):
-            del self._task_states[key]
         else:
+            self._task_states[key] = value
+
+    def __delitem__(self, key: asyncio.Task | Coroutine[Any, Any, Any], /) -> None:
+        if isinstance(key, Coroutine):
             del self._preliminary_task_states[key]
+        else:
+            del self._task_states[key]
 
     def __len__(self) -> int:
         return len(self._task_states) + len(self._preliminary_task_states)
 
-    def __iter__(self) -> Iterator[Awaitable[Any] | asyncio.Task]:
+    def __iter__(self) -> Iterator[Coroutine[Any, Any, Any] | asyncio.Task]:
         yield from self._task_states
         yield from self._preliminary_task_states
 


Test patch:
diff --git a/tests/test_taskgroups.py b/tests/test_taskgroups.py
index 1f53694..d20565c 100644
--- a/tests/test_taskgroups.py
+++ b/tests/test_taskgroups.py
@@ -11,7 +11,7 @@ from typing import Any, NoReturn, cast
 
 import pytest
 from exceptiongroup import catch
-from pytest import FixtureRequest
+from pytest import FixtureRequest, MonkeyPatch
 from pytest_mock import MockerFixture
 
 import anyio
@@ -1778,3 +1778,14 @@ async def test_eager_task_factory(request: FixtureRequest) -> None:
     async with create_task_group() as tg:
         tg.start_soon(sync_coro)
         tg.cancel_scope.cancel()
+
+
+@pytest.mark.parametrize("anyio_backend", ["asyncio"])
+async def test_patched_asyncio_task(monkeypatch: MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        asyncio,
+        "Task",
+        asyncio.tasks._PyTask,  # type: ignore[attr-defined]
+    )
+    async with create_task_group() as tg:
+        tg.start_soon(sleep, 0)