waldiez 0.5.9__py3-none-any.whl → 0.6.0__py3-none-any.whl

waldiez/running/base_runner.py

@@ -1,26 +1,20 @@
 # 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 json
 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 +25,31 @@ 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, dump_waldiez
 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]
+    # pylint: disable=line-too-long
+    from autogen.events import (  # type: ignore[import-untyped,import-not-found,unused-ignore]  # noqa: E501
+        BaseEvent,
+    )
+    from autogen.messages import (  # type: ignore[import-untyped,import-not-found,unused-ignore] # noqa: E501
+        BaseMessage,
+    )
 
 
-class WaldiezBaseRunner(WaldiezRunnerProtocol):
+# pylint: disable=too-many-public-methods
+class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
     """Base runner for Waldiez.
 
     Initialization parameters:
@@ -54,28 +57,29 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         - output_path: Path to save the output file.
         - uploads_root: Root directory for uploads.
         - structured_io: Whether to use structured I/O.
-        - 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
     _output_path: str | Path | None
     _uploads_root: str | Path | None
     _dot_env_path: str | Path | None
-    _skip_patch_io: bool
     _running: bool
+    _is_async: bool
+    _waldiez_file: Path
+    _input: Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
+    _print: Callable[..., None]
+    _send: Callable[[Union["BaseEvent", "BaseMessage"]], None]
 
     def __init__(
         self,
@@ -83,30 +87,59 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         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()
+        waldiez_file = kwargs.get("waldiez_file", "")
+        if isinstance(waldiez_file, str) and waldiez_file:
+            waldiez_file_path = Path(waldiez_file).resolve()
+        elif isinstance(waldiez_file, Path):
+            waldiez_file_path = waldiez_file.resolve()
+        else:
+            waldiez_file_path = dump_waldiez(waldiez, output_path=output_path)
+        if not waldiez_file_path or not waldiez_file_path.is_file():
+            raise ValueError("Could not resolve a waldiez file path")
+        WaldiezBaseRunner._waldiez_file = waldiez_file_path
+
+    @staticmethod
+    def print(*args: Any, **kwargs: Any) -> None:
+        """Print a message to the console.
+
+        Parameters
+        ----------
+        *args : Any
+            Positional arguments to print.
+        **kwargs : Any
+            Keyword arguments to print.
+        """
+        if len(args) == 1 and isinstance(args[0], dict):
+            arg = json.dumps(args[0], default=str, ensure_ascii=False)
+            WaldiezBaseRunner._print(arg, **kwargs)
+        else:
+            WaldiezBaseRunner._print(*args, **kwargs)
 
     def is_running(self) -> bool:
         """Check if the workflow is currently running.
@@ -116,65 +149,119 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         bool
             True if the workflow is running, False otherwise.
         """
-        return WaldiezBaseRunner._running
+        with self._running_lock:
+            return WaldiezBaseRunner._running
 
-    def wait_for_completion(self, timeout: Optional[float] = None) -> bool:
-        """Wait for non-blocking execution to complete.
+    @staticmethod
+    def get_input_function() -> (
+        Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
+    ):
+        """Get the input function for user interaction.
 
-        This is a base implementation that subclasses can override if needed.
+        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
+
+    @staticmethod
+    async def a_get_user_input(
+        prompt: str, *, password: bool = False, **kwargs: Any
+    ) -> 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).
+        **kwargs : Any
+            Additional keyword arguments to pass to the input function.
 
         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,
+                    **kwargs,
+                )
+            except TypeError:
+                result = await input_function(prompt)  # type: ignore
+        else:
+            try:
+                result = input_function(prompt, password=password, **kwargs)
+            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,
+        **kwargs: Any,
+    ) -> 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).
+        **kwargs : Any
+            Additional keyword arguments to pass to the input function.
 
         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, **kwargs
+                )
+            except TypeError:
+                return syncify(input_function)(prompt)
+        try:
+            return str(input_function(prompt, password=password, **kwargs))
+        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,48 +337,11 @@ 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]],
         output_file: Path,
+        waldiez_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
@@ -303,20 +353,26 @@ 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,
+                waldiez_file=waldiez_file,
+                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(
         self,
         results: list[dict[str, Any]],
         output_file: Path,
+        waldiez_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
@@ -330,23 +386,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
             skip_timeline=skip_timeline,
+            waldiez_file=waldiez_file,
         )
 
-    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 +408,67 @@ 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"],
+        skip_send: bool = False,
+    ) -> 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.
+        skip_send : bool
+            Skip sending the event.
+        """
+        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)
+            elif not skip_send:
+                WaldiezBaseRunner._send(event)
+
+    @staticmethod
+    def process_event(
+        event: Union["BaseEvent", "BaseMessage"],
+        skip_send: bool = False,
+    ) -> None:
+        """Process an event or message synchronously.
+
+        Parameters
+        ----------
+        event : Union[BaseEvent, BaseMessage]
+            The event or message to process.
+        skip_send : bool
+            Skip sending the event.
         """
-        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)
+            elif not skip_send:
+                WaldiezBaseRunner._send(event)
 
     def before_run(
         self,
@@ -452,6 +518,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 +564,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 +573,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 +586,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 +604,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 +613,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 +630,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 +651,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 +701,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 +734,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 +751,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)
@@ -627,6 +762,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             results=results,
             output_file=output_file,
             uploads_root=uploads_root_path,
+            waldiez_file=WaldiezBaseRunner._waldiez_file,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
             skip_timeline=skip_timeline,
@@ -635,129 +771,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]],
@@ -787,6 +800,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         self._after_run(
             results=results,
             output_file=output_file,
+            waldiez_file=WaldiezBaseRunner._waldiez_file,
             uploads_root=uploads_root,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
@@ -826,39 +840,23 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
             skip_timeline=skip_timeline,
+            waldiez_file=WaldiezBaseRunner._waldiez_file,
         )
 
-    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."""
         return self._waldiez
 
+    @property
+    def waldiez_file(self) -> Path:
+        """Get the path to the waldiez file."""
+        return self._waldiez_file
+
     @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:
@@ -890,11 +888,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         """Get the uploads root path for the runner."""
         return WaldiezBaseRunner._uploads_root
 
-    @property
-    def skip_patch_io(self) -> bool:
-        """Check if the runner is skipping patching IO."""
-        return WaldiezBaseRunner._skip_patch_io
-
     @classmethod
     def load(
         cls,
@@ -952,6 +945,31 @@ 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 set_stop_requested(self) -> None:
+        """Set the stop requested flag."""
+        self._stop_requested.set()
+
     def __enter__(self) -> Self:
         """Enter the context manager."""
         return self
@@ -968,7 +986,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 +996,4 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
     ) -> None:
         """Exit the context manager asynchronously."""
         if self.is_running():
-            await self.a_stop()
+            self._stop_requested.set()