diff --git a/helpers/runtime.py b/helpers/runtime.py index 5751de5f7a..6b070c230b 100644 --- a/helpers/runtime.py +++ b/helpers/runtime.py @@ -10,7 +10,23 @@ import sys import nest_asyncio -nest_asyncio.apply() + +def _apply_nest_asyncio_if_supported() -> None: + """Apply nest_asyncio unless the running loop is uvloop.""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + nest_asyncio.apply() + return + + loop_module = type(loop).__module__ + if loop_module == "uvloop" or loop_module.startswith("uvloop."): + return + + nest_asyncio.apply(loop) + + +_apply_nest_asyncio_if_supported() T = TypeVar("T") R = TypeVar("R") diff --git a/helpers/task_scheduler.py b/helpers/task_scheduler.py index fada618e91..cdfc223fea 100644 --- a/helpers/task_scheduler.py +++ b/helpers/task_scheduler.py @@ -10,7 +10,29 @@ from typing import Any, Callable, Dict, Literal, Optional, Type, TypeVar, Union, cast, ClassVar import nest_asyncio -nest_asyncio.apply() + + +def _apply_nest_asyncio_if_supported() -> None: + """Apply nest_asyncio unless the running loop is uvloop. + + nest_asyncio can only patch standard asyncio loops. Agent Zero's WebUI can + import this module while Uvicorn is running on uvloop; attempting to patch + that loop raises ValueError and breaks WebSocket connection setup. + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + nest_asyncio.apply() + return + + loop_module = type(loop).__module__ + if loop_module == "uvloop" or loop_module.startswith("uvloop."): + return + + nest_asyncio.apply(loop) + + +_apply_nest_asyncio_if_supported() from crontab import CronTab from pydantic import BaseModel, Field, PrivateAttr diff --git a/tests/test_task_scheduler_timezone.py b/tests/test_task_scheduler_timezone.py index 3d55780b83..3fe6796242 100644 --- a/tests/test_task_scheduler_timezone.py +++ b/tests/test_task_scheduler_timezone.py @@ -2,12 +2,16 @@ from datetime import datetime, timezone from pathlib import Path import sys -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) +plugins_pkg = ModuleType("plugins") +plugins_pkg.__path__ = [str(PROJECT_ROOT / "plugins")] +sys.modules["plugins"] = plugins_pkg + from helpers import task_scheduler from helpers.task_scheduler import AdHocTask, ScheduledTask, TaskSchedule diff --git a/tests/test_task_scheduler_uvloop_import.py b/tests/test_task_scheduler_uvloop_import.py new file mode 100644 index 0000000000..3aed026fc9 --- /dev/null +++ b/tests/test_task_scheduler_uvloop_import.py @@ -0,0 +1,30 @@ +import asyncio +import importlib +from pathlib import Path +import sys +import types + +import pytest + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + + +def _bind_local_plugins_namespace() -> None: + plugins_pkg = types.ModuleType("plugins") + plugins_pkg.__path__ = [str(PROJECT_ROOT / "plugins")] + sys.modules["plugins"] = plugins_pkg + + +def test_task_scheduler_import_is_safe_inside_uvloop_event_loop(): + uvloop = pytest.importorskip("uvloop") + sys.modules.pop("helpers.task_scheduler", None) + _bind_local_plugins_namespace() + + async def import_task_scheduler() -> None: + importlib.import_module("helpers.task_scheduler") + + with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner: + runner.run(import_task_scheduler())