intuned-runtime 1.2.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

runtime/browser/launch_browser.py

@@ -1,31 +1,59 @@
 from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack
+from typing import AsyncContextManager
+from typing import overload
+from typing import TYPE_CHECKING
 
-from playwright.async_api import ProxySettings
+if TYPE_CHECKING:
+    from playwright.async_api import BrowserContext
+    from playwright.async_api import Page
+    from playwright.async_api import ProxySettings
 
 from runtime.env import get_browser_type
+from runtime.errors.run_api_errors import AutomationError
 
 from .launch_camoufox import launch_camoufox
 from .launch_chromium import launch_chromium
 
 
+@overload
+def launch_browser(
+    *,
+    cdp_address: str,
+) -> "AsyncContextManager[tuple['BrowserContext', 'Page']]": ...
+
+
+@overload
+def launch_browser(
+    proxy: "ProxySettings | None" = None,
+    headless: bool = False,
+    *,
+    cdp_port: int | None = None,
+) -> "AsyncContextManager[tuple['BrowserContext', 'Page']]": ...
+
+
 @asynccontextmanager
 async def launch_browser(
-    proxy: ProxySettings | None = None,
+    proxy: "ProxySettings | None" = None,
     headless: bool = False,
     *,
+    cdp_port: int | None = None,
     cdp_address: str | None = None,
 ):
     browser_type = get_browser_type()
-    match browser_type:
-        case "camoufox":
-            async with launch_camoufox(headless=headless, proxy=proxy) as (context, page):
-                try:
-                    yield context, page
-                finally:
-                    await context.close()
-        case "chromium" | _:
-            async with launch_chromium(headless=headless, cdp_address=cdp_address, proxy=proxy) as (context, page):
-                try:
-                    yield context, page
-                finally:
-                    await context.close()
+    async with AsyncExitStack() as stack:
+        match browser_type:
+            case "camoufox":
+                if cdp_address:
+                    raise AutomationError(ValueError("CDP address is not supported with Camoufox"))
+                if cdp_port:
+                    raise AutomationError(ValueError("CDP port is not supported with Camoufox"))
+                context, page = await stack.enter_async_context(launch_camoufox(headless=headless, proxy=proxy))
+            case "chromium" | _:
+                context, page = await stack.enter_async_context(
+                    launch_chromium(headless=headless, cdp_address=cdp_address, cdp_port=cdp_port, proxy=proxy)
+                )
+        try:
+            yield context, page
+        finally:
+            await context.close()

runtime/browser/launch_chromium.py

@@ -4,13 +4,20 @@ import logging
 import os
 from contextlib import asynccontextmanager
 from typing import Any
+from typing import TYPE_CHECKING
 
 import anyio
 
+from .helpers import get_local_cdp_address
 from .helpers import get_proxy_env
+from .helpers import wait_on_cdp_address
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from playwright.async_api import ProxySettings
+    from playwright.async_api import ViewportSize
+
 
 chromium_launch_args_to_ignore = [
     "--disable-field-trial-config",
@@ -78,21 +85,17 @@ default_user_agent = (
 async def launch_chromium(
     headless: bool = True,
     timeout: int = 10,
+    cdp_port: int | None = None,
     cdp_address: str | None = None,
+    *,
+    proxy: "ProxySettings | None" = None,
+    viewport: "ViewportSize | None" = None,
     **kwargs: Any,
 ):
     from playwright.async_api import async_playwright
     from playwright.async_api import Browser
 
-    extra_args = [
-        "--no-first-run",
-        "--disable-sync",
-        "--disable-translate",
-        "--disable-features=TranslateUI",
-        "--disable-features=NetworkService",
-        "--lang=en",
-        "--disable-blink-features=AutomationControlled",
-    ]
+    extra_args: list[str] = []
     async with async_playwright() as playwright:
         if cdp_address is not None:
             browser: Browser = await playwright.chromium.connect_over_cdp(cdp_address)
@@ -104,29 +107,25 @@ async def launch_chromium(
                 user_preferences_dir,
                 dir_to_clean,
             ) = await create_user_dir_with_preferences()
-            if kwargs.get("proxy") is None:
-                proxy_env = get_proxy_env()
-            else:
-                proxy_env = kwargs.get("proxy")
+            proxy = proxy or get_proxy_env()
             # Remove proxy from kwargs if it exists
-            kwargs.pop("proxy", None)
-            viewport = kwargs.get("viewport", {"width": 1280, "height": 800})
-            kwargs.pop("viewport", None)
+            viewport = viewport or {"width": 1280, "height": 800}
 
-            if headless:
-                chromium_launch_args_to_ignore.append("--headless")
-                extra_args.append("--headless=new")
+            if cdp_port:
+                extra_args.append(f"--remote-debugging-port={cdp_port}")
 
             context = await playwright.chromium.launch_persistent_context(
                 os.fspath(user_preferences_dir),
                 headless=headless,
                 viewport=viewport,
-                proxy=proxy_env,
-                # ignore_default_args=chromium_launch_args_to_ignore,
+                proxy=proxy,
                 user_agent=os.environ.get("USER_AGENT", default_user_agent),
-                # args=extra_args,
+                args=extra_args,
                 **kwargs,
             )
+            if cdp_port:
+                await wait_on_cdp_address(get_local_cdp_address(cdp_port))
+
         context.set_default_timeout(timeout * 1000)
 
         async def remove_dir_after_close(*_: Any, **__: Any) -> None:
@@ -145,6 +144,7 @@ async def launch_chromium(
             await process.wait()
 
         context.once("close", remove_dir_after_close)
+
         yield context, context.pages[0]
 
 
@@ -173,6 +173,10 @@ async def dangerous_launch_chromium(
         browser: Browser = await playwright.chromium.connect_over_cdp(cdp_url)
         browser.on("disconnected", lambda _: logging.info("Browser Session disconnected"))
         context = browser.contexts[0]
+        # set view port for the already existing pages and any new pages
+        for page in context.pages:
+            await page.set_viewport_size(kwargs.get("viewport", {"width": 1280, "height": 800}))
+        context.on("page", lambda page: page.set_viewport_size(kwargs.get("viewport", {"width": 1280, "height": 800})))
         user_preferences_dir = None
         dir_to_clean = None
     else:

runtime/constants.py

@@ -0,0 +1 @@
+api_key_header_name = "x-api-key"

runtime/context/context.py

@@ -21,6 +21,7 @@ class IntunedContext(BaseModel):
     extended_payloads: list[Payload] = Field(default_factory=lambda: list[Payload]())
     run_context: IntunedRunContext | None = None
     get_auth_session_parameters: Callable[[], Awaitable[dict[str, Any]]] | None = None
+    store: dict[str, Any] = Field(default_factory=lambda: dict())
 
     _token: Token["IntunedContext"] | None = None
 

runtime/env.py

@@ -6,7 +6,7 @@ def get_workspace_id():
 
 
 def get_project_id():
-    return os.environ.get("INTUNED_INTEGRATION_ID")
+    return os.environ.get("INTUNED_INTEGRATION_ID", os.environ.get("INTUNED_PROJECT_ID"))
 
 
 def get_functions_domain():
@@ -15,3 +15,17 @@ def get_functions_domain():
 
 def get_browser_type():
     return os.environ.get("BROWSER_TYPE")
+
+
+def get_api_key():
+    return os.environ.get(api_key_env_var_key)
+
+
+def get_is_running_in_cli():
+    return os.environ.get(cli_env_var_key) == "true"
+
+
+cli_env_var_key = "INTUNED_CLI"
+workspace_env_var_key = "INTUNED_WORKSPACE_ID"
+project_env_var_key = "INTUNED_PROJECT_ID"
+api_key_env_var_key = "INTUNED_API_KEY"

runtime/errors/run_api_errors.py

@@ -58,6 +58,14 @@ class NoAutomationInApiError(RunApiError):
         )
 
 
+class InvalidAPIError(RunApiError):
+    def __init__(self, message: str):
+        super().__init__(
+            message,
+            "InvalidAPIError",
+        )
+
+
 class AutomationNotCoroutineError(RunApiError):
     def __init__(self, module: str, automation_function_name: str = "automation"):
         super().__init__(

runtime/helpers/__init__.py

@@ -1,5 +1,6 @@
+from .attempt_store import attempt_store
 from .extend_payload import extend_payload
 from .extend_timeout import extend_timeout
 from .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/helpers/attempt_store.py

@@ -0,0 +1,14 @@
+from typing import Any
+
+from runtime.context.context import IntunedContext
+
+
+class Store:
+    def get(self, key: str, default: Any = None) -> Any:
+        return IntunedContext.current().store.get(key, default)
+
+    def set(self, key: str, value: Any) -> None:
+        IntunedContext.current().store[key] = value
+
+
+attempt_store = Store()

runtime/run/playwright_context.py

@@ -0,0 +1,147 @@
+import asyncio
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack
+from typing import Any
+from typing import AsyncContextManager
+from typing import overload
+from typing import TYPE_CHECKING
+from typing import TypedDict
+from typing import Unpack
+
+from ..browser import launch_browser
+from ..browser.helpers import get_local_cdp_address
+from ..errors.run_api_errors import AutomationError
+from ..errors.run_api_errors import InvalidAPIError
+from ..run.types import ImportFunction
+from .setup_context_hook import load_setup_context_hook
+from .setup_context_hook import setup_context_hook_function_name
+
+if TYPE_CHECKING:
+    from playwright.async_api import BrowserContext
+    from playwright.async_api import Page
+    from playwright.async_api import ProxySettings
+
+
+class _CommonKwargs(TypedDict):
+    import_function: "ImportFunction"
+    api_name: str
+    api_parameters: Any
+
+
+@overload
+def playwright_context(
+    *,
+    proxy: "ProxySettings | None" = None,
+    headless: bool = False,
+    **kwargs: Unpack[_CommonKwargs],
+) -> "AsyncContextManager[tuple[BrowserContext, Page]]": ...
+@overload
+def playwright_context(
+    *,
+    cdp_address: str | None = None,
+    **kwargs: Unpack[_CommonKwargs],
+) -> "AsyncContextManager[tuple[BrowserContext, Page]]": ...
+
+
+@asynccontextmanager
+async def playwright_context(
+    *,
+    proxy: "ProxySettings | None" = None,
+    headless: bool = False,
+    cdp_address: str | None = None,
+    **kwargs: Unpack[_CommonKwargs],
+) -> "AsyncGenerator[tuple[BrowserContext, Page], None]":
+    setup_context_hook = load_setup_context_hook(import_function=kwargs["import_function"])
+
+    if setup_context_hook is None:
+        if cdp_address:
+            async with launch_browser(
+                cdp_address=cdp_address,
+            ) as (context, page):
+                yield context, page
+        else:
+            async with launch_browser(
+                proxy=proxy,
+                headless=headless,
+            ) as (context, page):
+                yield context, page
+        return
+
+    if cdp_address is not None:
+        cdp_port = None
+        hook_cdp_url = cdp_address
+    else:
+        cdp_port = await get_random_free_port()
+        hook_cdp_url = get_local_cdp_address(cdp_port)
+
+    async with AsyncExitStack() as stack:
+        if cdp_address:
+            context, page = await stack.enter_async_context(
+                launch_browser(
+                    cdp_address=cdp_address,
+                )
+            )
+        else:
+            context, page = await stack.enter_async_context(
+                launch_browser(
+                    proxy=proxy,
+                    headless=headless,
+                    cdp_port=cdp_port,
+                )
+            )
+        try:
+            hook_result = await setup_context_hook(
+                api_name=kwargs["api_name"],
+                api_parameters=kwargs["api_parameters"],
+                cdp_url=hook_cdp_url,
+            )
+        except Exception as e:
+            raise AutomationError(e) from e
+
+        if hook_result is None:
+            yield context, page
+            return
+
+        new_context: "BrowserContext | None" = None  # noqa: UP037
+        try:
+            if not isinstance(hook_result, tuple):
+                new_context = hook_result
+                yield new_context, page
+                return
+
+            if len(hook_result) == 2:
+                new_context, new_page = hook_result
+
+                yield new_context, new_page or page
+                return
+
+            if len(hook_result) == 3:
+                new_context, new_page, cleanup = hook_result
+
+                try:
+                    yield new_context, new_page or page
+                    return
+                finally:
+                    try:
+                        await cleanup()
+                    except Exception as e:
+                        raise AutomationError(e) from e
+
+            raise InvalidAPIError(
+                f"{setup_context_hook_function_name} hook returned an invalid value. Return value must be one of: None, BrowserContext, (BrowserContext, Page | None), (BrowserContext, Page | None, Callable[..., Awaitable[None]])"
+            )
+        finally:
+            if new_context is not None and new_context != context:
+                await new_context.close()
+
+
+async def get_random_free_port() -> int:
+    def get_random_free_port_sync():
+        import socket
+
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+
+    return await asyncio.to_thread(get_random_free_port_sync)

runtime/run/playwright_tracing.py

@@ -0,0 +1,27 @@
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from playwright.async_api import BrowserContext
+
+from anyio import Path
+
+
+@asynccontextmanager
+async def playwright_tracing(
+    *,
+    context: "BrowserContext",
+    trace_path: str,
+    screenshots: bool = True,
+    snapshots: bool = True,
+    sources: bool = True,
+):
+    await context.tracing.start(screenshots=screenshots, snapshots=snapshots, sources=sources)
+    try:
+        yield
+    finally:
+        try:
+            await context.tracing.stop(path=trace_path)
+        except Exception as e:
+            print("Error stopping tracing:", e)
+            await Path(trace_path).unlink(missing_ok=True)

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")