intuned-runtime 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

runtime/run/run_api.py

@@ -1,21 +1,18 @@
 import asyncio
 import functools
 import json
-import os.path
+import sys
+from collections.abc import Callable
+from contextlib import AsyncExitStack
 from importlib import import_module
-from importlib.util import module_from_spec
-from importlib.util import spec_from_file_location
 from inspect import iscoroutinefunction
 from typing import Any
-from typing import Callable
-from typing import Coroutine
+from typing import Awaitable
+from typing import cast
 from typing import Optional
 from typing import Protocol
 
-from tenacity import retry
-from tenacity import retry_if_not_result
-from tenacity import RetryError
-from tenacity import stop_after_attempt
+from git import TYPE_CHECKING
 
 from runtime.browser.storage_state import get_storage_state
 from runtime.browser.storage_state import set_storage_state
@@ -23,6 +20,7 @@ from runtime.context import IntunedContext
 from runtime.errors.run_api_errors import InvalidSessionError
 from runtime.errors.run_api_errors import ResultTooBigError
 from runtime.errors.run_api_errors import RunApiError
+from runtime.run.playwright_context import playwright_context
 from runtime.types import RunAutomationSuccessResult
 from runtime.types.run_types import PayloadToAppend
 from runtime.types.run_types import RunApiParameters
@@ -31,8 +29,12 @@ from ..errors import ApiNotFoundError
 from ..errors import AutomationError
 from ..errors import AutomationNotCoroutineError
 from ..errors import NoAutomationInApiError
-from .playwright_constructs import get_production_playwright_constructs
+from .playwright_tracing import playwright_tracing
 from .pydantic_encoder import PydanticEncoder
+from .types import ImportFunction
+
+if TYPE_CHECKING:
+    from playwright.async_api import Page
 
 
 def get_object_size_in_bytes(obj: Any) -> int:
@@ -49,31 +51,20 @@ def import_function_from_api_dir(
     file_path: str,
     base_dir: Optional[str] = None,
     automation_function_name: str | None = None,
-) -> Callable[..., Coroutine[Any, Any, Any]]:
+) -> Callable[..., Awaitable[Any]]:
     module_path = file_path.replace("/", ".")
 
     def _import_module():
-        if base_dir is None:
-            return import_module(module_path)
-        else:
-            file_location = os.path.join(base_dir, f"{file_path}.py")
-            if not os.path.exists(file_location):
-                raise ApiNotFoundError(module_path)
-            spec = spec_from_file_location(os.path.basename(file_path), file_location)
-            if spec is None:
-                raise ApiNotFoundError(module_path)
-            module = module_from_spec(spec)
-            if spec.loader is None:
-                raise ApiNotFoundError(module_path)
-            spec.loader.exec_module(module)
-            return module
+        if base_dir is not None:
+            sys.path.insert(0, base_dir)
+        return import_module(module_path)
 
     try:
         module = _import_module()
 
     except ModuleNotFoundError as e:
         # if the top-level module does not exist, it is a 404
-        if e.name == module_path:
+        if e.name == module_path or e.name == module_path.split(".", 1)[0]:
             raise ApiNotFoundError(module_path) from e
 
         # otherwise, it is an import error inside the user code
@@ -113,8 +104,10 @@ def import_function_from_api_dir(
     return automation_coroutine
 
 
-class ImportFunction(Protocol):
-    def __call__(self, file_path: str, name: Optional[str] = None, /) -> Callable[..., Coroutine[Any, Any, Any]]: ...
+class AutomationFunction(Protocol):
+    def __call__(
+        self, page: "Page", parameters: dict[Any, Any] | None = None, *args: Any, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
 
 async def run_api(
@@ -124,30 +117,50 @@ async def run_api(
 ) -> RunAutomationSuccessResult:
     from playwright.async_api import ProxySettings
 
-    trace_started: bool = False
+    import_function = import_function or (
+        lambda file_path, name=None: import_function_from_api_dir(file_path=file_path, automation_function_name=name)
+    )
 
-    headless = False
-    proxy = None
-    cdp_address = None
+    async with AsyncExitStack() as stack:
+        _initialize_playwright_context = functools.partial(
+            playwright_context,
+            import_function=import_function,
+            api_name=parameters.automation_function.name,
+            api_parameters=parameters.automation_function.params,
+        )
+        if parameters.run_options.environment == "standalone":
+            proxy_config = parameters.run_options.proxy
+            if proxy_config is not None:
+                proxy = ProxySettings(
+                    **proxy_config.model_dump(by_alias=True),
+                )
+            else:
+                proxy = None
 
-    if parameters.run_options.environment == "standalone":
-        headless = parameters.run_options.headless
-        proxy_config = parameters.run_options.proxy
-        if proxy_config is not None:
-            proxy = ProxySettings(
-                **proxy_config.model_dump(by_alias=True),
+            context, page = await stack.enter_async_context(
+                _initialize_playwright_context(
+                    proxy=proxy,
+                    headless=parameters.run_options.headless,
+                )
+            )
+
+        else:
+            context, page = await stack.enter_async_context(
+                _initialize_playwright_context(
+                    cdp_address=parameters.run_options.cdp_address,
+                )
             )
-    else:
-        cdp_address = parameters.run_options.cdp_address
 
-    async with get_production_playwright_constructs(
-        headless=headless,
-        proxy=proxy,
-        cdp_address=cdp_address,
-    ) as (context, page):
-        if parameters.tracing.enabled:
-            await context.tracing.start(screenshots=True, snapshots=True, sources=True)
-            trace_started = True
+        if parameters.tracing.enabled is True:
+            await stack.enter_async_context(
+                playwright_tracing(
+                    context=context,
+                    trace_path=parameters.tracing.file_path,
+                    screenshots=True,
+                    snapshots=True,
+                    sources=True,
+                )
+            )
 
         if parameters.auth is not None and parameters.auth.session.type == "state":
             if parameters.auth.session.state is None:
@@ -157,40 +170,16 @@ async def run_api(
                 context=context,
                 state=state,
             )
-        import_function = import_function or (
-            lambda file_path, name=None: import_function_from_api_dir(
-                file_path=file_path, automation_function_name=name
-            )
-        )
 
         async def _run_automation():
             try:
-                automation_coroutine = import_function(parameters.automation_function.name)
-
-                if parameters.auth is not None and parameters.auth.run_check:
-                    retry_configs = retry(
-                        stop=stop_after_attempt(2),
-                        retry=retry_if_not_result(lambda result: result is True),
-                        reraise=True,
-                    )
-
-                    check_fn = import_function("auth-sessions/check")
-
-                    check_fn_with_retries = retry_configs(check_fn)
-                    try:
-                        check_result = await check_fn_with_retries(page)
-                    except RetryError:
-                        check_result = False
-                    if type(check_result) is not bool:
-                        raise AutomationError(TypeError("Check function must return a boolean"))
-                    if not check_result:
-                        raise InvalidSessionError()
-
-                automation_coroutine_with_page = functools.partial(automation_coroutine, page)
+                automation_function = cast(AutomationFunction, import_function(parameters.automation_function.name))
+
+                automation_function = functools.partial(automation_function, page)
                 if parameters.automation_function.params is None:
-                    automation_result = await automation_coroutine_with_page()
+                    automation_result = await automation_function()
                 else:
-                    automation_result = await automation_coroutine_with_page(parameters.automation_function.params)
+                    automation_result = await automation_function(parameters.automation_function.params)
                 try:
                     automation_result = json.loads(json.dumps(automation_result, cls=PydanticEncoder))
                 except TypeError as e:
@@ -228,16 +217,13 @@ async def run_api(
                 # Get all public attributes of the exception
                 raise AutomationError(e) from e
 
-        automation_task = None
+        automation_task = asyncio.create_task(_run_automation())
         try:
-            automation_task = asyncio.create_task(_run_automation())
-
             # Shield will make the CancelledError get thrown directly here instead of inside `automation_task`
-            result = await asyncio.shield(automation_task)
-            return result
+            return await asyncio.shield(automation_task)
         except asyncio.CancelledError:
             # Manually cancel the automation task
-            if automation_task and not automation_task.done():
+            if not automation_task.done():
                 automation_task.cancel()
                 try:
                     # Wait for the automation task to be cancelled for a brief moment
@@ -245,11 +231,5 @@ async def run_api(
                 except (asyncio.CancelledError, asyncio.TimeoutError):
                     pass
             raise  # Re-raise the cancellation
-        finally:
-            if parameters.tracing.enabled is True and trace_started:
-                try:
-                    await context.tracing.stop(path=parameters.tracing.file_path)
-                except Exception as e:
-                    print("Error stopping tracing:", e)
-                    os.remove(parameters.tracing.file_path)
-            await context.close()
+
+    raise RuntimeError("Unreachable code path")

runtime/run/setup_context_hook.py

@@ -0,0 +1,40 @@
+from collections.abc import Callable
+from typing import Any
+from typing import Awaitable
+from typing import cast
+from typing import Protocol
+from typing import TYPE_CHECKING
+from typing import Union
+
+from runtime.errors.run_api_errors import ApiNotFoundError
+
+from .types import ImportFunction
+
+if TYPE_CHECKING:
+    from playwright.async_api import BrowserContext
+    from playwright.async_api import Page
+
+
+setup_context_hook_path = "hooks/setup_context"
+setup_context_hook_function_name = "setup_context"
+
+SetupContextHookReturn = Union[
+    "None",
+    "BrowserContext",
+    "tuple[BrowserContext, Page | None]",
+    "tuple[BrowserContext, Page | None, Callable[..., Awaitable[None]]]",
+]
+
+
+class SetupContextHook(Protocol):
+    def __call__(self, *, cdp_url: str, api_name: str, api_parameters: Any) -> Awaitable[SetupContextHookReturn]: ...
+
+
+def load_setup_context_hook(*, import_function: ImportFunction):
+    try:
+        setup_context_hook = cast(
+            SetupContextHook, import_function(setup_context_hook_path, setup_context_hook_function_name)
+        )
+        return setup_context_hook
+    except ApiNotFoundError:
+        return None

runtime/run/types.py

@@ -0,0 +1,14 @@
+from collections.abc import Callable
+from typing import Any
+from typing import Awaitable
+from typing import Optional
+from typing import Protocol
+
+from git import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    pass
+
+
+class ImportFunction(Protocol):
+    def __call__(self, file_path: str, name: Optional[str] = None, /) -> Callable[..., Awaitable[Any]]: ...

runtime/types/run_types.py

@@ -126,7 +126,6 @@ class TracingDisabled(CamelBaseModel):
 
 class Auth(CamelBaseModel):
     session: RunApiSession
-    run_check: bool = False
 
 
 class IntunedRunContext(CamelBaseModel):

runtime_helpers/__init__.py

@@ -1,5 +1,6 @@
 from runtime.helpers import extend_payload
 from runtime.helpers import extend_timeout
+from runtime.helpers.attempt_store import attempt_store
 from runtime.helpers.get_auth_session_parameters import get_auth_session_parameters
 
-__all__ = ["extend_payload", "extend_timeout", "get_auth_session_parameters"]
+__all__ = ["extend_payload", "extend_timeout", "get_auth_session_parameters", "attempt_store"]

runtime/run/playwright_constructs.py

@@ -1,20 +0,0 @@
-from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from playwright.async_api import ProxySettings
-from ..browser import launch_browser
-
-
-@asynccontextmanager
-async def get_production_playwright_constructs(
-    proxy: "ProxySettings | None" = None,
-    headless: bool = False,
-    *,
-    cdp_address: str | None = None,
-):
-    async with launch_browser(headless=headless, cdp_address=cdp_address, proxy=proxy) as (context, page):
-        try:
-            yield context, page
-        finally:
-            await context.close()