waldiez 0.5.8__py3-none-any.whl → 0.5.10__py3-none-any.whl

waldiez/running/base_runner.py

@@ -1,26 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
-# pylint: disable=too-many-instance-attributes,unused-argument
-# pylint: disable=too-many-arguments,too-many-positional-arguments
-# pylint: disable=too-many-public-methods,too-many-locals
-
+# pyright: reportUnknownMemberType=false, reportAttributeAccessIssue=false
+# pyright: reportUnknownArgumentType=false
 """Base runner for Waldiez workflows."""
 
+import importlib.util
+import inspect
 import shutil
 import sys
 import tempfile
 import threading
 from pathlib import Path
-from types import TracebackType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Coroutine,
-    Optional,
-    Type,
-)
+from types import ModuleType, TracebackType
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, Type, Union
 
 from aiofiles.os import wrap
 from anyio.from_thread import start_blocking_portal
@@ -31,22 +24,26 @@ from waldiez.logger import WaldiezLogger, get_logger
 from waldiez.models import Waldiez
 
 from .environment import refresh_environment, reset_env_vars, set_env_vars
+from .exceptions import StopRunningException
 from .post_run import after_run
-from .pre_run import (
-    a_install_requirements,
-    install_requirements,
-)
+from .pre_run import RequirementsMixin
 from .protocol import WaldiezRunnerProtocol
 from .utils import (
     a_chdir,
     chdir,
+    input_async,
+    input_sync,
+    is_async_callable,
+    syncify,
 )
 
 if TYPE_CHECKING:
     from autogen.events import BaseEvent  # type: ignore[import-untyped]
+    from autogen.messages import BaseMessage  # type: ignore[import-untyped]
 
 
-class WaldiezBaseRunner(WaldiezRunnerProtocol):
+# pylint: disable=too-many-public-methods
+class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
     """Base runner for Waldiez.
 
     Initialization parameters:
@@ -57,17 +54,15 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         - skip_patch_io: Whether to skip patching I/O functions.
         - dot_env: Path to a .env file for environment variables.
 
-    Methods to override:
+    Methods to possibly override:
+        - prepare: Prepare the environment and paths for running the flow.
         - _before_run: Actions to perform before running the flow.
+        - a_prepare: Async version of the prepare method.
         - _a_before_run: Async actions to perform before running the flow.
         - _run: Actual implementation of the run logic.
         - _a_run: Async implementation of the run logic.
         - _after_run: Actions to perform after running the flow.
         - _a_after_run: Async actions to perform after running the flow.
-        - _start: Implementation of non-blocking start logic.
-        - _a_start: Async implementation of non-blocking start logic.
-        - _stop: Actions to perform when stopping the flow.
-        - _a_stop: Async actions to perform when stopping the flow.
     """
 
     _structured_io: bool
@@ -76,6 +71,10 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
     _dot_env_path: str | Path | None
     _skip_patch_io: bool
     _running: bool
+    _is_async: bool
+    _input: Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
+    _print: Callable[..., None]
+    _send: Callable[[Union["BaseEvent", "BaseMessage"]], None]
 
     def __init__(
         self,
@@ -85,28 +84,32 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io: bool,
         skip_patch_io: bool = False,
         dot_env: str | Path | None = None,
+        **kwargs: Any,
     ) -> None:
         """Initialize the Waldiez manager."""
-        self._waldiez = waldiez
         WaldiezBaseRunner._running = False
         WaldiezBaseRunner._structured_io = structured_io
         WaldiezBaseRunner._output_path = output_path
         WaldiezBaseRunner._uploads_root = uploads_root
         WaldiezBaseRunner._skip_patch_io = skip_patch_io
         WaldiezBaseRunner._dot_env_path = dot_env
+        WaldiezBaseRunner._input = input
+        WaldiezBaseRunner._print = print
+        WaldiezBaseRunner._send = print
+        WaldiezBaseRunner._is_async = waldiez.is_async
+        self._waldiez = waldiez
         self._called_install_requirements = False
         self._exporter = WaldiezExporter(waldiez)
         self._stop_requested = threading.Event()
-        self._logger = get_logger()
-        self._print: Callable[..., None] = print
-        self._send: Callable[["BaseEvent"], None] = self._print
-        self._input: (
-            Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
-        ) = input
         self._last_results: list[dict[str, Any]] = []
         self._last_exception: Exception | None = None
-        self._execution_complete_event = threading.Event()
-        self._execution_thread: Optional[threading.Thread] = None
+        self._running_lock = threading.Lock()
+        self._loaded_module: ModuleType | None = None
+        logger = kwargs.get("logger")
+        if isinstance(logger, WaldiezLogger):
+            self._logger = logger
+        else:
+            self._logger = get_logger()
 
     def is_running(self) -> bool:
         """Check if the workflow is currently running.
@@ -116,65 +119,109 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         bool
             True if the workflow is running, False otherwise.
         """
-        return WaldiezBaseRunner._running
+        with self._running_lock:
+            return WaldiezBaseRunner._running
+
+    @staticmethod
+    def get_input_function() -> (
+        Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
+    ):
+        """Get the input function for user interaction.
 
-    def wait_for_completion(self, timeout: Optional[float] = None) -> bool:
-        """Wait for non-blocking execution to complete.
+        Returns
+        -------
+        Callable[[str, bool], str]
+            A function that takes a prompt and a password flag,
+            returning user input.
+        """
+        if hasattr(WaldiezBaseRunner, "_input") and callable(
+            WaldiezBaseRunner._input
+        ):
+            return WaldiezBaseRunner._input
+        if WaldiezBaseRunner._is_async:
+            return input_async
+        return input_sync
 
-        This is a base implementation that subclasses can override if needed.
+    @staticmethod
+    async def a_get_user_input(prompt: str, *, password: bool = False) -> str:
+        """Get user input with an optional password prompt.
 
         Parameters
         ----------
-        timeout: float | None
-            The maximum time to wait for completion, in seconds.
-            If None, wait indefinitely.
+        prompt : str
+            The prompt to display to the user.
+        password : bool, optional
+            If True, the input will be hidden (default is False).
 
         Returns
         -------
-        bool
-            True if the workflow completed successfully, False if it timed out.
+        str
+            The user input.
         """
-        # Default implementation - subclasses can override
-        if not self.is_running():
-            return True
-
-        if self._execution_thread and self._execution_thread.is_alive():
-            self._execution_thread.join(timeout)
-            return not self._execution_thread.is_alive()
+        input_function = WaldiezBaseRunner.get_input_function()
+        if is_async_callable(input_function):
+            try:
+                result = await input_function(  # type: ignore
+                    prompt,
+                    password=password,
+                )
+            except TypeError:
+                result = await input_function(prompt)  # type: ignore
+        else:
+            try:
+                result = input_function(prompt, password=password)
+            except TypeError:
+                result = input_function(prompt)
+        return result  # pyright: ignore
 
-        return self._execution_complete_event.wait(timeout or 0)
+    @staticmethod
+    def get_user_input(
+        prompt: str,
+        *,
+        password: bool = False,
+    ) -> str:
+        """Get user input with an optional password prompt.
 
-    def get_execution_stats(self) -> dict[str, Any]:
-        """Get basic execution statistics.
+        Parameters
+        ----------
+        prompt : str
+            The prompt to display to the user.
+        password : bool, optional
+            If True, the input will be hidden (default is False).
 
         Returns
         -------
-        dict[str, Any]
-            A dictionary containing execution statistics.
-        - is_running: Whether the runner is currently running.
-        - has_thread: Whether there is an execution thread.
-        - thread_alive: Whether the execution thread is alive.
-        - has_error: Whether there was an error during execution.
+        str
+            The user input.
         """
-        return {
-            "is_running": self.is_running(),
-            "has_thread": self._execution_thread is not None,
-            "thread_alive": (
-                self._execution_thread.is_alive()
-                if self._execution_thread
-                else False
-            ),
-            "has_error": self._last_exception is not None,
-        }
-
-    # Helper for subclasses
-    def _signal_completion(self) -> None:
-        """Signal that execution has completed."""
-        self._execution_complete_event.set()
-
-    def _reset_completion_state(self) -> None:
-        """Reset completion state for new execution."""
-        self._execution_complete_event.clear()
+        input_function = WaldiezBaseRunner.get_input_function()
+        if inspect.iscoroutinefunction(input_function):
+            try:
+                return syncify(input_function)(prompt, password=password)
+            except TypeError:
+                return syncify(input_function)(prompt)
+        try:
+            return str(input_function(prompt, password=password))
+        except TypeError:
+            return str(input_function(prompt))
+
+    def _load_module(self, output_file: Path, temp_dir: Path) -> ModuleType:
+        """Load the module from the waldiez file."""
+        file_name = output_file.name
+        module_name = file_name.replace(".py", "")
+        spec = importlib.util.spec_from_file_location(
+            module_name, temp_dir / file_name
+        )
+        if not spec or not spec.loader:
+            raise ImportError("Could not import the flow")
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        if not hasattr(module, "main"):
+            raise ImportError(
+                "The waldiez file does not contain a main() function"
+            )
+        self._loaded_module = module
+        return module
 
     def _before_run(
         self,
@@ -250,44 +297,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             "The _a_run method must be implemented in the subclass."
         )
 
-    def _start(
-        self,
-        temp_dir: Path,
-        output_file: Path,
-        uploads_root: Path | None,
-        skip_mmd: bool,
-        skip_timeline: bool,
-    ) -> None:
-        """Start running the Waldiez flow in a non-blocking way."""
-        raise NotImplementedError(
-            "The _start method must be implemented in the subclass."
-        )
-
-    async def _a_start(
-        self,
-        temp_dir: Path,
-        output_file: Path,
-        uploads_root: Path | None,
-        skip_mmd: bool,
-        skip_timeline: bool,
-    ) -> None:
-        """Start running the Waldiez flow in a non-blocking way asynchronously.
-
-        Parameters
-        ----------
-        temp_dir : Path
-            The path to the temporary directory created for the run.
-        output_file : Path
-            The path to the output file.
-        uploads_root : Path | None
-            The root path for uploads, if any.
-        skip_mmd : bool
-            Whether to skip generating the mermaid diagram.
-        """
-        raise NotImplementedError(
-            "The _a_start method must be implemented in the subclass."
-        )
-
     def _after_run(
         self,
         results: list[dict[str, Any]],
@@ -303,14 +312,18 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
 
         # Reset stop flag for next run
         self._stop_requested.clear()
-        after_run(
-            temp_dir=temp_dir,
-            output_file=output_file,
-            flow_name=self.waldiez.name,
-            uploads_root=uploads_root,
-            skip_mmd=skip_mmd,
-            skip_timeline=skip_timeline,
-        )
+        # pylint: disable=broad-exception-caught
+        try:
+            after_run(
+                temp_dir=temp_dir,
+                output_file=output_file,
+                flow_name=self._waldiez.name,
+                uploads_root=uploads_root,
+                skip_mmd=skip_mmd,
+                skip_timeline=skip_timeline,
+            )
+        except BaseException as exc:  # pragma: no cover
+            self.log.warning("Error occurred during after_run: %s", exc)
         self.log.info("Cleanup completed")
 
     async def _a_after_run(
@@ -332,21 +345,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             skip_timeline=skip_timeline,
         )
 
-    def _stop(self) -> None:
-        """Actions to perform when stopping the flow."""
-        raise NotImplementedError(
-            "The _stop method must be implemented in the subclass."
-        )
-
-    async def _a_stop(self) -> None:
-        """Asynchronously perform actions when stopping the flow."""
-        raise NotImplementedError(
-            "The _a_stop method must be implemented in the subclass."
-        )
-
-    # ===================================================================
-    # HELPER METHODS
-    # ===================================================================
     @staticmethod
     def _prepare_paths(
         output_path: str | Path | None = None,
@@ -366,43 +364,57 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_file: Path = Path(WaldiezBaseRunner._output_path)
         return output_file, uploads_root_path
 
-    def gather_requirements(self) -> set[str]:
-        """Gather extra requirements to install before running the flow.
+    @staticmethod
+    async def a_process_event(event: Union["BaseEvent", "BaseMessage"]) -> None:
+        """Process an event or message asynchronously.
 
-        Returns
-        -------
-        set[str]
-            A set of requirements that are not already installed and do not
-            include 'waldiez' in their name.
+        Parameters
+        ----------
+        event : Union[BaseEvent, BaseMessage]
+            The event or message to process.
+        """
+        if hasattr(event, "type"):  # pragma: no branch
+            if event.type == "input_request":
+                prompt = getattr(
+                    event, "prompt", getattr(event.content, "prompt", "> ")
+                )
+                password = getattr(
+                    event,
+                    "password",
+                    getattr(event.content, "password", False),
+                )
+                user_input = await WaldiezBaseRunner.a_get_user_input(
+                    prompt, password=password
+                )
+                await event.content.respond(user_input)
+            else:
+                WaldiezBaseRunner._send(event)
+
+    @staticmethod
+    def process_event(event: Union["BaseEvent", "BaseMessage"]) -> None:
+        """Process an event or message synchronously.
+
+        Parameters
+        ----------
+        event : Union[BaseEvent, BaseMessage]
+            The event or message to process.
         """
-        extra_requirements = {
-            req
-            for req in self.waldiez.requirements
-            if req not in sys.modules and "waldiez" not in req
-        }
-        if "python-dotenv" not in extra_requirements:
-            extra_requirements.add("python-dotenv")
-        return extra_requirements
-
-    def install_requirements(self) -> None:
-        """Install the requirements for the flow."""
-        if not self._called_install_requirements:
-            self._called_install_requirements = True
-            extra_requirements = self.gather_requirements()
-            if extra_requirements:
-                install_requirements(extra_requirements)
-
-    async def a_install_requirements(self) -> None:
-        """Install the requirements for the flow asynchronously."""
-        if not self._called_install_requirements:
-            self._called_install_requirements = True
-            extra_requirements = self.gather_requirements()
-            if extra_requirements:
-                await a_install_requirements(extra_requirements)
-
-    # ===================================================================
-    # PUBLIC PROTOCOL IMPLEMENTATION
-    # ===================================================================
+        if hasattr(event, "type"):  # pragma: no branch
+            if event.type == "input_request":
+                prompt = getattr(
+                    event, "prompt", getattr(event.content, "prompt", "> ")
+                )
+                password = getattr(
+                    event,
+                    "password",
+                    getattr(event.content, "password", False),
+                )
+                user_input = WaldiezBaseRunner.get_user_input(
+                    prompt, password=password
+                )
+                event.content.respond(user_input)
+            else:
+                WaldiezBaseRunner._send(event)
 
     def before_run(
         self,
@@ -452,6 +464,42 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             uploads_root=uploads_root,
         )
 
+    def prepare(
+        self,
+        output_path: str | Path | None,
+        uploads_root: str | Path | None,
+    ) -> tuple[Path, Path, Path | None]:
+        """Prepare the paths and environment for running the flow.
+
+        Parameters
+        ----------
+        output_path : str | Path | None
+            The output path for the flow, by default None.
+        uploads_root : str | Path | None
+            The root path for uploads, by default None.
+
+        Returns
+        -------
+        tuple[Path, Path, Path | None]
+            A tuple containing:
+            - The path to the output file.
+            - The path to the temporary directory created for the run.
+            - The root path for uploads, if specified, otherwise None.
+        """
+        output_file, uploads_root_path = self._prepare_paths(
+            output_path=output_path,
+            uploads_root=uploads_root,
+        )
+        temp_dir = self.before_run(
+            output_file=output_file,
+            uploads_root=uploads_root_path,
+        )
+        self.install_requirements()
+        refresh_environment()
+        return temp_dir, output_file, uploads_root_path
+
+    # noinspection PyProtocol
+    # pylint: disable=too-many-locals,unused-argument
     def run(
         self,
         output_path: str | Path | None = None,
@@ -462,7 +510,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         dot_env: str | Path | None = None,
         **kwargs: Any,
     ) -> list[dict[str, Any]]:
-        """Run the Waldiez flow in blocking mode.
+        """Run the Waldiez flow.
 
         Parameters
         ----------
@@ -471,10 +519,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         uploads_root : str | Path | None
             The runtime uploads root, by default None.
         structured_io : bool
-            Whether to use structured IO instead of the default 'input/print',
-            by default False.
+            Whether to use structured IO instead of the default 'input/print'.
         skip_mmd : bool
-            Whether to skip generating the mermaid diagram, by default False.
+            Whether to skip generating the mermaid diagram.
         skip_timeline : bool
             Whether to skip generating the timeline JSON.
         dot_env : str | Path | None
@@ -485,12 +532,15 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         Returns
         -------
         list[dict[str, Any]]
-            The result of the run.
+            The results of the run.
 
         Raises
         ------
         RuntimeError
-            If the runner is already running.
+            If the runner is already running, the workflow is not async,
+            or an error occurs during the run.
+        StopRunningException
+            If the run is stopped by the user.
         """
         if dot_env is not None:
             resolved = Path(dot_env).resolve()
@@ -500,7 +550,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
             raise RuntimeError("Workflow already running")
-        if self.waldiez.is_async:
+        if self.is_async:
             with start_blocking_portal(backend="asyncio") as portal:
                 return portal.call(
                     self.a_run,
@@ -509,19 +559,13 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                     structured_io,
                     skip_mmd,
                 )
-        output_file, uploads_root_path = self._prepare_paths(
+        temp_dir, output_file, uploads_root_path = self.prepare(
             output_path=output_path,
             uploads_root=uploads_root,
         )
-        temp_dir = self.before_run(
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-        )
-        self.install_requirements()
-        refresh_environment()
         WaldiezBaseRunner._running = True
-        results: list[dict[str, Any]] = []
-        old_env_vars = set_env_vars(self.waldiez.get_flow_env_vars())
+        results: list[dict[str, Any]]
+        old_env_vars = set_env_vars(self._waldiez.get_flow_env_vars())
         try:
             with chdir(to=temp_dir):
                 sys.path.insert(0, str(temp_dir))
@@ -532,6 +576,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                     skip_mmd=skip_mmd,
                     skip_timeline=skip_timeline,
                 )
+        except (SystemExit, StopRunningException, KeyboardInterrupt) as exc:
+            raise StopRunningException(StopRunningException.reason) from exc
+        except BaseException as exc:  # pylint: disable=broad-exception-caught
+            self.log.error("Error occurred while running workflow: %s", exc)
+            results = [{"error": str(exc)}]
         finally:
             WaldiezBaseRunner._running = False
             reset_env_vars(old_env_vars)
@@ -548,7 +597,39 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             sys.path.pop(0)
         return results
 
+    async def a_prepare(
+        self,
+        output_path: str | Path | None,
+        uploads_root: str | Path | None,
+    ) -> tuple[Path, Path, Path | None]:
+        """Prepare the paths for the async run.
+
+        Parameters
+        ----------
+        output_path : str | Path | None
+            The output path, by default None.
+        uploads_root : str | Path | None
+            The uploads root path, by default None.
+
+        Returns
+        -------
+        tuple[Path, Path, Path | None]
+            The temporary directory, output file, and uploads root path.
+        """
+        output_file, uploads_root_path = self._prepare_paths(
+            output_path=output_path,
+            uploads_root=uploads_root,
+        )
+        temp_dir = await self._a_before_run(
+            output_file=output_file,
+            uploads_root=uploads_root_path,
+        )
+        await self.a_install_requirements()
+        refresh_environment()
+        return temp_dir, output_file, uploads_root_path
+
     # noinspection DuplicatedCode
+    # noinspection PyProtocol
     async def a_run(
         self,
         output_path: str | Path | None = None,
@@ -566,28 +647,30 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path : str | Path | None
             The output path, by default None.
         uploads_root : str | Path | None
-            The runtime uploads root, by default None.
+            The runtime uploads root.
         structured_io : bool
-            Whether to use structured IO instead of the default 'input/print',
-            by default False.
+            Whether to use structured IO instead of the default 'input/print'.
         skip_mmd : bool
-            Whether to skip generating the mermaid diagram, by default False.
+            Whether to skip generating the mermaid diagram.
         skip_timeline : bool
-            Whether to skip generating the timeline JSON, by default False.
+            Whether to skip generating the timeline JSON.
         dot_env : str | Path | None
             The path to the .env file, if any.
         **kwargs : Any
-            Additional keyword arguments for the run method.
+            Additional keyword arguments for the a_run method.
 
         Returns
         -------
         list[dict[str, Any]]
-            The result of the run.
+            The results of the run.
 
         Raises
         ------
         RuntimeError
-            If the runner is already running.
+            If the runner is already running, the workflow is not async
+            or an error occurs during the run.
+        StopRunningException
+            If the run is stopped by the user.
         """
         if dot_env is not None:
             resolved = Path(dot_env).resolve()
@@ -597,19 +680,13 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
             raise RuntimeError("Workflow already running")
-        output_file, uploads_root_path = self._prepare_paths(
+        temp_dir, output_file, uploads_root_path = await self.a_prepare(
             output_path=output_path,
             uploads_root=uploads_root,
         )
-        temp_dir = await self._a_before_run(
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-        )
-        await self.a_install_requirements()
-        refresh_environment()
         WaldiezBaseRunner._running = True
-        results: list[dict[str, Any]] = []
-        old_env_vars = set_env_vars(self.waldiez.get_flow_env_vars())
+        results: list[dict[str, Any]]
+        old_env_vars = set_env_vars(self._waldiez.get_flow_env_vars())
         try:
             async with a_chdir(to=temp_dir):
                 sys.path.insert(0, str(temp_dir))
@@ -620,6 +697,10 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                     skip_mmd=skip_mmd,
                     skip_timeline=skip_timeline,
                 )
+        except (SystemExit, StopRunningException, KeyboardInterrupt) as exc:
+            raise StopRunningException(StopRunningException.reason) from exc
+        except BaseException as exc:  # pylint: disable=broad-exception-caught
+            results = [{"error": str(exc)}]
         finally:
             WaldiezBaseRunner._running = False
             reset_env_vars(old_env_vars)
@@ -635,129 +716,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             sys.path.pop(0)
         return results
 
-    def start(
-        self,
-        output_path: str | Path | None,
-        uploads_root: str | Path | None,
-        structured_io: bool | None = None,
-        skip_mmd: bool = False,
-        skip_timeline: bool = False,
-        dot_env: str | Path | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Start running the Waldiez flow in a non-blocking way.
-
-        Parameters
-        ----------
-        output_path : str | Path | None
-            The output path.
-        uploads_root : str | Path | None
-            The runtime uploads root.
-        structured_io : bool | None
-            Whether to use structured IO instead of the default 'input/print'.
-        skip_mmd : bool
-            Whether to skip generating the mermaid diagram, by default False.
-        skip_timeline : bool
-            Whether to skip generating the timeline JSON, by default False.
-        dot_env : str | Path | None
-            The path to the .env file, if any.
-        **kwargs : Any
-            Additional keyword arguments for the start method.
-
-        Raises
-        ------
-        RuntimeError
-            If the runner is already running.
-        """
-        if dot_env is not None:
-            resolved = Path(dot_env).resolve()
-            if resolved.is_file():
-                WaldiezBaseRunner._dot_env_path = resolved
-        if structured_io is not None:
-            WaldiezBaseRunner._structured_io = structured_io
-        if self.is_running():
-            raise RuntimeError("Workflow already running")
-        output_file, uploads_root_path = self._prepare_paths(
-            output_path=output_path,
-            uploads_root=uploads_root,
-        )
-        temp_dir = self.before_run(
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-        )
-        self.install_requirements()
-        refresh_environment()
-        WaldiezBaseRunner._running = True
-        self._start(
-            temp_dir=temp_dir,
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-            skip_mmd=skip_mmd,
-            skip_timeline=skip_timeline,
-        )
-
-    # noinspection DuplicatedCode
-    async def a_start(
-        self,
-        output_path: str | Path | None,
-        uploads_root: str | Path | None,
-        structured_io: bool | None = None,
-        skip_mmd: bool = False,
-        skip_timeline: bool = False,
-        dot_env: str | Path | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Asynchronously start running the Waldiez flow in a non-blocking way.
-
-        Parameters
-        ----------
-        output_path : str | Path | None
-            The output path.
-        uploads_root : str | Path | None
-            The runtime uploads root.
-        structured_io : bool | None = None
-            Whether to use structured IO instead of the default 'input/print'.
-        skip_mmd : bool = False
-            Whether to skip generating the mermaid diagram, by default False.
-        skip_timeline : bool = False
-            Whether to skip generating the timeline JSON, by default False.
-        dot_env : str | Path | None = None
-            The path to the .env file, if any.
-        **kwargs : Any
-            Additional keyword arguments for the start method.
-
-        Raises
-        ------
-        RuntimeError
-            If the runner is already running.
-        """
-        if dot_env is not None:
-            resolved = Path(dot_env).resolve()
-            if resolved.is_file():
-                WaldiezBaseRunner._dot_env_path = resolved
-        if structured_io is not None:
-            WaldiezBaseRunner._structured_io = structured_io
-        if self.is_running():
-            raise RuntimeError("Workflow already running")
-        output_file, uploads_root_path = self._prepare_paths(
-            output_path=output_path,
-            uploads_root=uploads_root,
-        )
-        temp_dir = await self._a_before_run(
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-        )
-        await self.a_install_requirements()
-        refresh_environment()
-        WaldiezBaseRunner._running = True
-        await self._a_start(
-            temp_dir=temp_dir,
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-            skip_mmd=skip_mmd,
-            skip_timeline=skip_timeline,
-        )
-
     def after_run(
         self,
         results: list[dict[str, Any]],
@@ -828,28 +786,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             skip_timeline=skip_timeline,
         )
 
-    def stop(self) -> None:
-        """Stop the runner if it is running."""
-        if not self.is_running():
-            return
-        try:
-            self._stop()
-        finally:
-            WaldiezBaseRunner._running = False
-
-    async def a_stop(self) -> None:
-        """Asynchronously stop the runner if it is running."""
-        if not self.is_running():
-            return
-        try:
-            await self._a_stop()
-        finally:
-            WaldiezBaseRunner._running = False
-
-    # ===================================================================
-    # PROPERTIES AND CONTEXT MANAGERS
-    # ===================================================================
-
     @property
     def waldiez(self) -> Waldiez:
         """Get the Waldiez instance."""
@@ -858,7 +794,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
     @property
     def is_async(self) -> bool:
         """Check if the workflow is async."""
-        return self.waldiez.is_async
+        return self._waldiez.is_async
 
     @property
     def running(self) -> bool:
@@ -952,6 +888,27 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             dot_env=dot_env,
         )
 
+    def stop(self) -> None:
+        """Stop the workflow execution.
+
+        This method sets the stop flag that will be checked by the event
+        handlers and other parts of the workflow execution to gracefully
+        terminate the workflow execution.
+        Note: Stopping will occur at the "next" AutoGen event, not immediately.
+        """
+        self.log.info("Stop requested - setting stop flag")
+        self._stop_requested.set()
+
+    def is_stop_requested(self) -> bool:
+        """Check if a stop has been requested.
+
+        Returns
+        -------
+        bool
+            True if stop has been requested, False otherwise.
+        """
+        return self._stop_requested.is_set()
+
     def __enter__(self) -> Self:
         """Enter the context manager."""
         return self
@@ -968,7 +925,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
     ) -> None:
         """Exit the context manager."""
         if self.is_running():
-            self.stop()
+            self._stop_requested.set()
 
     async def __aexit__(
         self,
@@ -978,4 +935,4 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
     ) -> None:
         """Exit the context manager asynchronously."""
         if self.is_running():
-            await self.a_stop()
+            self._stop_requested.set()