intuned-runtime 1.3.0rc0__py3-none-any.whl → 1.3.2__py3-none-any.whl

intuned_cli/controller/api.py

@@ -1,27 +1,29 @@
 import json
-import os
 from typing import Any
+from typing import Unpack
 
 from anyio import Path
 from pydantic import BaseModel
 
 from intuned_cli.controller.authsession import execute_run_validate_auth_session_cli
 from intuned_cli.controller.authsession import load_auth_session_instance
+from intuned_cli.types import BaseExecuteCommandOptionsWithoutTrace
 from intuned_cli.utils.api_helpers import assert_api_file_exists
+from intuned_cli.utils.browser import get_cli_run_options
 from intuned_cli.utils.console import console
 from intuned_cli.utils.error import CLIError
 from intuned_cli.utils.error import log_automation_error
 from intuned_cli.utils.get_auth_parameters import register_get_auth_session_parameters
+from intuned_cli.utils.import_function import get_cli_import_function
 from intuned_cli.utils.timeout import extendable_timeout
+from intuned_cli.utils.traces import cli_trace
 from runtime.errors.run_api_errors import AutomationError
-from runtime.run.run_api import import_function_from_api_dir
 from runtime.run.run_api import run_api
 from runtime.types.run_types import Auth
 from runtime.types.run_types import AutomationFunction
 from runtime.types.run_types import PayloadToAppend
 from runtime.types.run_types import ProxyConfig
 from runtime.types.run_types import RunApiParameters
-from runtime.types.run_types import StandaloneRunOptions
 from runtime.types.run_types import StateSession
 from runtime.types.run_types import StorageState
 
@@ -40,9 +42,8 @@ async def execute_run_api_cli(
     retries: int,
     auth_session: AuthSessionInput | None = None,
     output_file: str | None = None,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    trace: bool,
+    **kwargs: Unpack[BaseExecuteCommandOptionsWithoutTrace],
 ):
     """
     Execute API with retries and optional auth session validation.
@@ -63,9 +64,8 @@ async def execute_run_api_cli(
                 auto_recreate=auth_session.auto_recreate,
                 check_retries=auth_session.check_retries,
                 create_retries=auth_session.create_retries,
-                headless=headless,
-                timeout=timeout,
-                proxy=proxy,
+                trace=trace,
+                **kwargs,
             )
 
         try:
@@ -73,9 +73,8 @@ async def execute_run_api_cli(
                 api_name=api_name,
                 parameters=input_data,
                 auth=auth_session_instance,
-                headless=headless,
-                timeout=timeout,
-                proxy=proxy,
+                trace_id=f"{api_name}-attempt-{i + 1}" if trace else None,
+                **kwargs,
             )
 
             return await handle_api_result(
@@ -101,9 +100,8 @@ async def execute_attempt_api_cli(
     input_data: Any | None,
     auth_session_id: str | None = None,
     output_file: str | None = None,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    trace: bool,
+    **kwargs: Unpack[BaseExecuteCommandOptionsWithoutTrace],
 ):
     """
     Execute a single API attempt with optional auth session.
@@ -122,9 +120,8 @@ async def execute_attempt_api_cli(
         api_name=api_name,
         parameters=input_data,
         auth=auth_session_instance,
-        headless=headless,
-        timeout=timeout,
-        proxy=proxy,
+        trace_id=f"{api_name}-attempt" if trace else None,
+        **kwargs,
     )
 
     return await handle_api_result(
@@ -186,26 +183,30 @@ async def attempt_api(
     api_name: str,
     parameters: Any | None,
     auth: StorageState | None = None,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    trace_id: str | None = None,
+    **kwargs: Unpack[BaseExecuteCommandOptionsWithoutTrace],
 ):
-    cwd = await Path().resolve()
-    async with extendable_timeout(timeout):
-        result = await run_api(
-            RunApiParameters(
-                automation_function=AutomationFunction(name=f"api/{api_name}", params=parameters),
-                auth=Auth(
-                    session=StateSession(state=auth),
-                )
-                if auth
-                else None,
-                run_options=StandaloneRunOptions(
-                    headless=headless, proxy=ProxyConfig.parse_from_str(proxy) if proxy else None
+    timeout = kwargs.get("timeout")
+    headless = kwargs.get("headless")
+    proxy = kwargs.get("proxy")
+    keep_browser_open = kwargs.get("keep_browser_open")
+    with cli_trace(trace_id) as tracing:
+        async with extendable_timeout(timeout):
+            result = await run_api(
+                RunApiParameters(
+                    automation_function=AutomationFunction(name=f"api/{api_name}", params=parameters),
+                    auth=Auth(
+                        session=StateSession(state=auth),
+                    )
+                    if auth
+                    else None,
+                    run_options=await get_cli_run_options(
+                        headless=headless,
+                        proxy=ProxyConfig.parse_from_str(proxy) if proxy else None,
+                        keep_browser_open=keep_browser_open,
+                    ),
+                    tracing=tracing,
                 ),
-            ),
-            import_function=lambda file_path, name=None: import_function_from_api_dir(
-                file_path=file_path, automation_function_name=name, base_dir=os.fspath(cwd)
-            ),
-        )
-    return result.result, result.payload_to_append
+                import_function=await get_cli_import_function(),
+            )
+        return result.result, result.payload_to_append

intuned_cli/controller/authsession.py

@@ -1,31 +1,41 @@
+import asyncio
 import datetime
 import json
 import time
 from typing import Any
 from typing import Literal
+from typing import TYPE_CHECKING
+from typing import Unpack
 
 from anyio import Path
 from pydantic import BaseModel
 from pydantic import Field
 from pydantic import ValidationError
 
+from intuned_cli.types import BaseExecuteCommandOptions
+from intuned_cli.types import BaseExecuteCommandOptionsWithoutTrace
 from intuned_cli.utils.api_helpers import assert_api_file_exists
+from intuned_cli.utils.browser import get_cli_run_options
 from intuned_cli.utils.console import console
 from intuned_cli.utils.error import CLIError
 from intuned_cli.utils.error import log_automation_error
 from intuned_cli.utils.get_auth_parameters import register_get_auth_session_parameters
 from intuned_cli.utils.import_function import get_cli_import_function
 from intuned_cli.utils.timeout import extendable_timeout
+from intuned_cli.utils.traces import cli_trace
+from runtime.browser.storage_state import get_storage_state
 from runtime.errors.run_api_errors import AutomationError
 from runtime.run.run_api import run_api
 from runtime.types.run_types import Auth
 from runtime.types.run_types import AutomationFunction
 from runtime.types.run_types import ProxyConfig
 from runtime.types.run_types import RunApiParameters
-from runtime.types.run_types import StandaloneRunOptions
 from runtime.types.run_types import StateSession
 from runtime.types.run_types import StorageState
 
+if TYPE_CHECKING:
+    from playwright.async_api import ProxySettings
+
 auth_session_instances_dirname = "auth-session-instances"
 
 
@@ -45,9 +55,7 @@ async def execute_run_validate_auth_session_cli(
     auto_recreate: bool,
     check_retries: int,
     create_retries: int,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    **kwargs: Unpack[BaseExecuteCommandOptions],
 ) -> StorageState:
     """
     Validate auth session with optional auto-recreation.
@@ -64,13 +72,13 @@ async def execute_run_validate_auth_session_cli(
     check_result = await run_check_with_retries(
         auth=instance,
         retries=check_retries,
-        headless=headless,
-        timeout=timeout,
-        proxy=proxy,
+        **kwargs,
     )
 
     if not check_result:
         if not auto_recreate:
+            if metadata and metadata.auth_session_type == "MANUAL":
+                raise CLIError("Auth session validation failed")
             raise CLIError("Auto recreate is disabled, please provide a new auth session or update it manually")
 
         if metadata and metadata.auth_session_type == "MANUAL":
@@ -86,19 +94,15 @@ async def execute_run_validate_auth_session_cli(
             auth_session_id=id,
             auth_session_input=auth_session_input,
             retries=create_retries,
-            headless=headless,
-            timeout=timeout,
-            proxy=proxy,
             metadata=metadata,
+            **kwargs,
         )
 
         # Rerun check after refresh
         check_result = await run_check_with_retries(
             auth=instance,
             retries=check_retries,
-            headless=headless,
-            timeout=timeout,
-            proxy=proxy,
+            **kwargs,
         )
         if not check_result:
             raise CLIError("Failed to re-create auth session")
@@ -114,16 +118,14 @@ async def execute_run_create_auth_session_cli(
     check_retries: int,
     create_retries: int,
     log: bool = True,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
     metadata: AuthSessionMetadata | None = None,
+    **kwargs: Unpack[BaseExecuteCommandOptions],
 ):
     """
     Create a new auth session.
     """
     if id is None:
-        id = f"auth-session-{int(time.time() * 1000)}"
+        id = generate_auth_session_id()
 
     if log:
         console.print(f"[bold]Creating auth session with id [cyan]{id}[/cyan][/bold]")
@@ -135,18 +137,14 @@ async def execute_run_create_auth_session_cli(
         auth_session_id=id,
         auth_session_input=input_data,
         retries=create_retries,
-        headless=headless,
-        timeout=timeout,
-        proxy=proxy,
         metadata=metadata,
+        **kwargs,
     )
 
     check_result = await run_check_with_retries(
         auth=instance,
         retries=check_retries,
-        headless=headless,
-        timeout=timeout,
-        proxy=proxy,
+        **kwargs,
     )
     if not check_result:
         raise CLIError("Failed to create auth session")
@@ -161,9 +159,7 @@ async def execute_run_update_auth_session_cli(
     input_data: Any | None = None,
     check_retries: int,
     create_retries: int,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    **kwargs: Unpack[BaseExecuteCommandOptions],
 ):
     """
     Update an existing auth session.
@@ -186,10 +182,8 @@ async def execute_run_update_auth_session_cli(
         check_retries=check_retries,
         create_retries=create_retries,
         log=False,
-        headless=headless,
-        timeout=timeout,
-        proxy=proxy,
         metadata=metadata,
+        **kwargs,
     )
 
     console.print("[bold][green]Auth session updated successfully[/green][/bold]")
@@ -199,35 +193,25 @@ async def execute_attempt_create_auth_session_cli(
     *,
     id: str | None = None,
     input_data: Any,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    **kwargs: Unpack[BaseExecuteCommandOptions],
 ):
     """
     Execute a single attempt to create auth session.
     """
     if id is None:
-        id = f"auth-session-attempt-{int(time.time() * 1000)}"
+        id = generate_auth_session_id(is_attempt=True)
 
     console.print(f"[bold]Executing create auth session attempt with id [cyan]{id}[/cyan][/bold]")
     await assert_api_file_exists("auth-sessions", "create")
     await run_create_with_retries(
         auth_session_id=id,
         auth_session_input=input_data,
-        retries=1,
-        headless=headless,
-        timeout=timeout,
-        proxy=proxy,
+        retries=None,
+        **kwargs,
     )
 
 
-async def execute_attempt_check_auth_session_cli(
-    *,
-    id: str,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
-):
+async def execute_attempt_check_auth_session_cli(*, id: str, **kwargs: Unpack[BaseExecuteCommandOptions]):
     """
     Execute a single attempt to check auth session.
     """
@@ -240,10 +224,8 @@ async def execute_attempt_check_auth_session_cli(
 
     check_result = await run_check_with_retries(
         auth=instance,
-        retries=1,
-        headless=headless,
-        timeout=timeout,
-        proxy=proxy,
+        retries=None,
+        **kwargs,
     )
 
     if not check_result:
@@ -253,90 +235,105 @@ async def execute_attempt_check_auth_session_cli(
 
 
 async def run_check(
-    *,
-    auth: StorageState,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    *, auth: StorageState, trace_id: str | None = None, **kwargs: Unpack[BaseExecuteCommandOptionsWithoutTrace]
 ) -> bool:
     """
     Run auth session check.
     """
-    async with extendable_timeout(timeout):
-        result = await run_api(
-            RunApiParameters(
-                automation_function=AutomationFunction(
-                    name="auth-sessions/check",
-                    params=None,
-                ),
-                run_options=StandaloneRunOptions(
-                    headless=headless, proxy=ProxyConfig.parse_from_str(proxy) if proxy else None
-                ),
-                auth=Auth(
-                    session=StateSession(
-                        state=auth,
+    timeout = kwargs.get("timeout")
+    headless = kwargs.get("headless")
+    proxy = kwargs.get("proxy")
+    keep_browser_open = kwargs.get("keep_browser_open")
+    with cli_trace(trace_id) as tracing:
+        async with extendable_timeout(timeout):
+            result = await run_api(
+                RunApiParameters(
+                    automation_function=AutomationFunction(
+                        name="auth-sessions/check",
+                        params=None,
+                    ),
+                    run_options=await get_cli_run_options(
+                        headless=headless,
+                        proxy=ProxyConfig.parse_from_str(proxy) if proxy else None,
+                        keep_browser_open=keep_browser_open,
                     ),
+                    auth=Auth(
+                        session=StateSession(
+                            state=auth,
+                        ),
+                    ),
+                    tracing=tracing,
                 ),
-            ),
-            import_function=await get_cli_import_function(),
-        )
+                import_function=await get_cli_import_function(),
+            )
 
-        if not result.result:
-            return False
+            if not result.result:
+                return False
 
-        return bool(result.result)
+            return bool(result.result)
 
 
 async def run_create(
     *,
     auth_session_input: dict[str, Any],
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    trace_id: str | None = None,
+    **kwargs: Unpack[BaseExecuteCommandOptionsWithoutTrace],
 ) -> StorageState:
     """
     Run auth session create.
     """
-    async with extendable_timeout(timeout):
-        result = await run_api(
-            RunApiParameters(
-                automation_function=AutomationFunction(
-                    name="auth-sessions/create",
-                    params=auth_session_input,
-                ),
-                run_options=StandaloneRunOptions(
-                    headless=headless, proxy=ProxyConfig.parse_from_str(proxy) if proxy else None
+
+    timeout = kwargs.get("timeout")
+    headless = kwargs.get("headless")
+    proxy = kwargs.get("proxy")
+    keep_browser_open = kwargs.get("keep_browser_open")
+    with cli_trace(trace_id) as tracing:
+        async with extendable_timeout(timeout):
+            result = await run_api(
+                RunApiParameters(
+                    automation_function=AutomationFunction(
+                        name="auth-sessions/create",
+                        params=auth_session_input,
+                    ),
+                    run_options=await get_cli_run_options(
+                        headless=headless,
+                        proxy=ProxyConfig.parse_from_str(proxy) if proxy else None,
+                        keep_browser_open=keep_browser_open,
+                    ),
+                    retrieve_session=True,
+                    tracing=tracing,
                 ),
-                retrieve_session=True,
-            ),
-            import_function=await get_cli_import_function(),
-        )
-        if not result.session:
-            raise Exception("Auth session create did not return a session")
-        return result.session
+                import_function=await get_cli_import_function(),
+            )
+            if not result.session:
+                raise Exception("Auth session create did not return a session")
+            return result.session
 
 
 async def run_check_with_retries(
     *,
     auth: StorageState,
-    retries: int,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    retries: int | None,
+    trace: bool,
+    **kwargs: Unpack[BaseExecuteCommandOptionsWithoutTrace],
 ) -> bool:
     """
     Run auth session check with retries.
     """
-    for i in range(retries):
+    for i in range(retries or 1):
         attempt_text = "" if i == 0 else f" [italic](Attempt {i + 1})[/italic]"
         console.print(f"\n[bold]Running [cyan]auth session check[/cyan]{attempt_text}...[/bold]\n")
 
         try:
+            trace_id: str | None = None
+            if trace:
+                trace_id = "authsession-check-attempt"
+                if retries is not None:
+                    trace_id += f"-{i+1}"
             check_result = await run_check(
                 auth=auth,
-                headless=headless,
-                timeout=timeout,
-                proxy=proxy,
+                trace_id=trace_id,
+                **kwargs,
             )
 
             if check_result:
@@ -354,27 +351,31 @@ async def run_create_with_retries(
     *,
     auth_session_id: str,
     auth_session_input: dict[str, Any],
-    retries: int,
-    headless: bool,
-    timeout: float,
-    proxy: str | None = None,
+    retries: int | None,
     metadata: AuthSessionMetadata | None = None,
+    trace: bool,
+    **kwargs: Unpack[BaseExecuteCommandOptionsWithoutTrace],
 ):
     """
     Run auth session create with retries.
     """
     new_auth_session_instance: StorageState | None = None
 
-    for i in range(retries):
+    for i in range(retries or 1):
         attempt_text = "" if i == 0 else f" [italic](Attempt {i + 1})[/italic]"
         console.print(f"\n[bold]Running [cyan]auth session create[/cyan]{attempt_text}...[/bold]\n")
 
         try:
+            trace_id: str | None = None
+            if trace:
+                trace_id = "authsession-create-attempt"
+                if retries is not None:
+                    trace_id += f"-{i+1}"
+
             new_auth_session_instance = await run_create(
                 auth_session_input=auth_session_input,
-                headless=headless,
-                timeout=timeout,
-                proxy=proxy,
+                trace_id=trace_id,
+                **kwargs,
             )
             console.print("[bold][green]Auth session create succeeded[/green][/bold]")
             break
@@ -424,7 +425,7 @@ async def load_auth_session_instance(auth_session_id: str) -> tuple[StorageState
 async def store_auth_session_instance(
     auth_session_instance: StorageState,
     auth_session_id: str,
-    auth_session_input: dict[str, Any],
+    auth_session_input: dict[str, Any] | None = None,
     metadata: AuthSessionMetadata | None = None,
 ):
     """
@@ -445,13 +446,115 @@ async def store_auth_session_instance(
         updatedAt=datetime.datetime.now().isoformat(),
         authSessionId=auth_session_id,
         authSessionInput=auth_session_input,
-        authSessionType="API",
+        authSessionType=metadata.auth_session_type if metadata else "API",
     )
     metadata_file_path = auth_session_path / "metadata.json"
 
-    metadata_json = json.dumps(metadata.model_dump(by_alias=True), indent=2)
+    metadata_json = json.dumps(metadata.model_dump(by_alias=True, exclude_none=True), indent=2)
     await metadata_file_path.write_text(metadata_json)
 
 
 def get_auth_session_path(auth_session_id: str):
     return Path(auth_session_instances_dirname) / auth_session_id
+
+
+async def execute_record_auth_session_cli(
+    *,
+    start_url: str,
+    finish_url: str,
+    id: str | None = None,
+    check_retries: int = 1,
+    **kwargs: Unpack[BaseExecuteCommandOptions],
+):
+    """
+    Record a new auth session using a browser.
+    """
+
+    from playwright.async_api import ProxySettings
+
+    if id is None:
+        id = generate_auth_session_id()
+
+    console.print(f"[bold]Recording auth session with id [cyan]{id}[/cyan][/bold]")
+    proxy = kwargs.get("proxy")
+
+    proxy_settings: ProxySettings | None = None
+    if proxy:
+        proxy_settings = ProxySettings(server=proxy)
+
+    try:
+        auth_session = await record_auth_session(
+            start_url=start_url,
+            finish_url=finish_url,
+            timeout=kwargs.get("timeout"),
+            proxy=proxy_settings,
+        )
+    except CLIError as e:
+        raise CLIError(f"Failed to record auth session: {e}") from e
+
+    await store_auth_session_instance(
+        auth_session_instance=auth_session,
+        auth_session_id=id,
+        metadata=AuthSessionMetadata(
+            createdAt=datetime.datetime.now().isoformat(),
+            updatedAt=datetime.datetime.now().isoformat(),
+            authSessionId=id,
+            authSessionInput=None,
+            authSessionType="MANUAL",
+            recorderStartUrl=start_url,
+            recorderEndUrl=finish_url,
+        ),
+    )
+
+    await execute_run_validate_auth_session_cli(
+        id=id,
+        create_retries=1,
+        check_retries=check_retries,
+        auto_recreate=False,
+        **kwargs,
+    )
+
+    console.print(f"[bold][green]Auth session [cyan]{id}[/cyan] recorded successfully[/green][/bold]")
+
+
+async def record_auth_session(
+    *,
+    start_url: str,
+    finish_url: str,
+    timeout: float = 300,
+    proxy: "ProxySettings | None" = None,
+):
+    from runtime.browser import launch_chromium
+
+    async with launch_chromium(
+        proxy=proxy,
+        headless=False,
+        app_mode_initial_url=start_url,
+    ) as (context, page):
+        if not page.url.startswith(start_url):
+            await page.goto(start_url)
+        console.print(f"[bold]Navigated to[/bold] [underline]{start_url}[/underline]")
+        console.print(f"[bold]Waiting for[/bold] [underline]{finish_url}[/underline]...")
+
+        try:
+            async with asyncio.timeout(timeout):
+                while True:
+                    if len(context.pages) == 0:
+                        raise CLIError("Browser was closed before reaching the finish URL")
+                    if context.pages[0].url.startswith(finish_url):
+                        break
+                    await asyncio.sleep(1)
+
+            console.print("[bold]Finish URL reached, capturing auth session...[/bold]")
+            await page.wait_for_load_state("load")
+            auth_session = await get_storage_state(context)
+            return auth_session
+        except asyncio.TimeoutError as e:
+            raise CLIError("Timeout waiting for finish URL") from e
+
+
+def generate_auth_session_id(is_attempt: bool = False) -> str:
+    timestamp = int(time.time() * 1000)
+    if is_attempt:
+        return f"auth-session-attempt-{timestamp}"
+    return f"auth-session-{timestamp}"