SDK: various small improvements (#4170)

skyvern/library/skyvern.py

@@ -3,10 +3,12 @@ import os
 from typing import Any, overload
 
 import httpx
+import structlog
 from dotenv import load_dotenv
 from playwright.async_api import Playwright, async_playwright
 
 from skyvern.client import AsyncSkyvern, BrowserSessionResponse, SkyvernEnvironment
+from skyvern.client.core import RequestOptions
 from skyvern.client.types.task_run_response import TaskRunResponse
 from skyvern.client.types.workflow_run_response import WorkflowRunResponse
 from skyvern.forge.sdk.api.llm.models import LLMConfig, LLMRouterConfig
@@ -16,6 +18,8 @@ from skyvern.library.skyvern_browser import SkyvernBrowser
 from skyvern.schemas.run_blocks import CredentialType
 from skyvern.schemas.runs import ProxyLocation, RunEngine, RunStatus
 
+LOG = structlog.get_logger()
+
 
 class Skyvern(AsyncSkyvern):
     """Main entry point for the Skyvern SDK.
@@ -234,6 +238,7 @@ class Skyvern(AsyncSkyvern):
         include_action_history_in_verification: bool | None = None,
         max_screenshot_scrolls: int | None = None,
         browser_address: str | None = None,
+        request_options: RequestOptions | None = None,
     ) -> TaskRunResponse:
         task_run = await super().run_task(
             prompt=prompt,
@@ -255,6 +260,7 @@ class Skyvern(AsyncSkyvern):
             include_action_history_in_verification=include_action_history_in_verification,
             max_screenshot_scrolls=max_screenshot_scrolls,
             browser_address=browser_address,
+            request_options=request_options,
         )
 
         if wait_for_completion:
@@ -287,6 +293,7 @@ class Skyvern(AsyncSkyvern):
         run_with: str | None = None,
         wait_for_completion: bool = False,
         timeout: float = DEFAULT_AGENT_TIMEOUT,
+        request_options: RequestOptions | None = None,
     ) -> WorkflowRunResponse:
         workflow_run = await super().run_workflow(
             workflow_id=workflow_id,
@@ -306,6 +313,7 @@ class Skyvern(AsyncSkyvern):
             browser_address=browser_address,
             ai_fallback=ai_fallback,
             run_with=run_with,
+            request_options=request_options,
         )
         if wait_for_completion:
             async with asyncio.timeout(timeout):
@@ -341,6 +349,7 @@ class Skyvern(AsyncSkyvern):
         azure_vault_totp_secret_key: str | None = None,
         wait_for_completion: bool = False,
         timeout: float = DEFAULT_AGENT_TIMEOUT,
+        request_options: RequestOptions | None = None,
     ) -> WorkflowRunResponse:
         workflow_run = await super().login(
             credential_type=credential_type,
@@ -363,6 +372,7 @@ class Skyvern(AsyncSkyvern):
             azure_vault_username_key=azure_vault_username_key,
             azure_vault_password_key=azure_vault_password_key,
             azure_vault_totp_secret_key=azure_vault_totp_secret_key,
+            request_options=request_options,
         )
         if wait_for_completion:
             async with asyncio.timeout(timeout):
@@ -423,27 +433,54 @@ class Skyvern(AsyncSkyvern):
         """
         self._ensure_cloud_environment()
         browser_session = await self.get_browser_session(browser_session_id)
+        LOG.info("Connecting to existing cloud browser session", browser_session_id=browser_session.browser_session_id)
         return await self._connect_to_cloud_browser_session(browser_session)
 
-    async def launch_cloud_browser(self) -> SkyvernBrowser:
+    async def launch_cloud_browser(
+        self,
+        *,
+        timeout: int | None = None,
+        proxy_location: ProxyLocation | None = None,
+    ) -> SkyvernBrowser:
         """Launch a new cloud-hosted browser session.
 
         This creates a new browser session in Skyvern's cloud infrastructure and connects to it.
 
+        Args:
+            timeout: Timeout in minutes for the session. Timeout is applied after the session is started.
+                Must be between 5 and 1440. Defaults to 60.
+            proxy_location: Geographic proxy location to route the browser traffic through.
+                This is only available in Skyvern Cloud.
+
         Returns:
             SkyvernBrowser: A browser instance connected to the new cloud session.
         """
         self._ensure_cloud_environment()
-        browser_session = await self.create_browser_session()
+        browser_session = await self.create_browser_session(
+            timeout=timeout,
+            proxy_location=proxy_location,
+        )
+        LOG.info("Launched new cloud browser session", browser_session_id=browser_session.browser_session_id)
         return await self._connect_to_cloud_browser_session(browser_session)
 
-    async def use_cloud_browser(self) -> SkyvernBrowser:
+    async def use_cloud_browser(
+        self,
+        *,
+        timeout: int | None = None,
+        proxy_location: ProxyLocation | None = None,
+    ) -> SkyvernBrowser:
         """Get or create a cloud browser session.
 
         This method attempts to reuse the most recent available cloud browser session.
         If no session exists, it creates a new one. This is useful for cost efficiency
         and session persistence.
 
+        Args:
+            timeout: Timeout in minutes for the session. Timeout is applied after the session is started.
+                Must be between 5 and 1440. Defaults to 60. Only used when creating a new session.
+            proxy_location: Geographic proxy location to route the browser traffic through.
+                This is only available in Skyvern Cloud. Only used when creating a new session.
+
         Returns:
             SkyvernBrowser: A browser instance connected to an existing or new cloud session.
         """
@@ -453,7 +490,15 @@ class Skyvern(AsyncSkyvern):
             (s for s in browser_sessions if s.runnable_id is None), key=lambda s: s.started_at, default=None
         )
         if browser_session is None:
-            browser_session = await self.create_browser_session()
+            LOG.info("No existing cloud browser session found, launching a new session")
+            browser_session = await self.create_browser_session(
+                timeout=timeout,
+                proxy_location=proxy_location,
+            )
+            LOG.info("Launched new cloud browser session", browser_session_id=browser_session.browser_session_id)
+        else:
+            LOG.info("Reusing existing cloud browser session", browser_session_id=browser_session.browser_session_id)
+
         return await self._connect_to_cloud_browser_session(browser_session)
 
     def _ensure_cloud_environment(self) -> None:

skyvern/library/skyvern_browser.py

@@ -105,3 +105,25 @@ class SkyvernBrowser(BrowserContext):
 
     async def _create_skyvern_page(self, page: Page) -> SkyvernBrowserPage:
         return SkyvernBrowserPage(self, page)
+
+    async def close(self, **kwargs: Any) -> None:
+        """Close the browser and optionally close the browser session.
+
+        This method closes the browser context. If the browser is associated with a
+        cloud browser session (has a browser_session_id), it will also close the
+        browser session via the API, marking it as completed.
+
+        Args:
+            **kwargs: Arguments passed to the underlying BrowserContext.close() method.
+
+        Example:
+            ```python
+            browser = await skyvern.launch_cloud_browser()
+            # ... use the browser ...
+            await browser.close()  # Closes both browser and cloud session
+            ```
+        """
+        await self._browser_context.close(**kwargs)
+
+        if self._browser_session_id:
+            await self._skyvern.close_browser_session(self._browser_session_id)

skyvern/library/skyvern_browser_page.py

@@ -5,6 +5,7 @@ import structlog
 from playwright.async_api import Page
 
 from skyvern.client import GetRunResponse
+from skyvern.client.core import RequestOptions
 from skyvern.client.types.workflow_run_response import WorkflowRunResponse
 from skyvern.core.script_generations.skyvern_page import SkyvernPage
 from skyvern.library.constants import DEFAULT_AGENT_HEARTBEAT_INTERVAL, DEFAULT_AGENT_TIMEOUT
@@ -85,7 +86,9 @@ class SkyvernPageRun:
             browser_session_id=self._browser.browser_session_id,
             browser_address=self._browser.browser_address,
             user_agent=user_agent,
+            request_options=RequestOptions(additional_headers={"X-User-Agent": "skyvern-sdk"}),
         )
+        LOG.info("AI task is running, this may take a while", run_id=task_run.run_id)
 
         task_run = await self._wait_for_run_completion(task_run.run_id, timeout)
         return TaskRunResponse.model_validate(task_run.model_dump())
@@ -110,7 +113,7 @@ class SkyvernPageRun:
         """Run a login task in the context of this page and wait for it to finish.
 
         Args:
-            credential_type: Type of credential store to use (e.g., bitwarden, onepassword).
+            credential_type: Type of credential store to use (e.g., skyvern, bitwarden, onepassword).
             url: URL to navigate to for login. If not provided, uses the current page URL.
             credential_id: ID of the credential to use.
             bitwarden_collection_id: Bitwarden collection ID containing the credentials.
@@ -128,7 +131,7 @@ class SkyvernPageRun:
             WorkflowRunResponse containing the login workflow execution results.
         """
 
-        LOG.info("AI login", prompt=prompt)
+        LOG.info("Starting AI login workflow", credential_type=credential_type)
 
         workflow_run = await self._browser.skyvern.login(
             credential_type=credential_type,
@@ -145,7 +148,9 @@ class SkyvernPageRun:
             browser_session_id=self._browser.browser_session_id,
             browser_address=self._browser.browser_address,
             extra_http_headers=extra_http_headers,
+            request_options=RequestOptions(additional_headers={"X-User-Agent": "skyvern-sdk"}),
         )
+        LOG.info("AI login workflow is running, this may take a while", run_id=workflow_run.run_id)
 
         workflow_run = await self._wait_for_run_completion(workflow_run.run_id, timeout)
         return WorkflowRunResponse.model_validate(workflow_run.model_dump())
@@ -177,7 +182,7 @@ class SkyvernPageRun:
             WorkflowRunResponse containing the workflow execution results.
         """
 
-        LOG.info("AI run workflow", workflow_id=workflow_id)
+        LOG.info("Starting AI workflow", workflow_id=workflow_id)
 
         workflow_run = await self._browser.skyvern.run_workflow(
             workflow_id=workflow_id,
@@ -189,7 +194,9 @@ class SkyvernPageRun:
             totp_identifier=totp_identifier,
             browser_session_id=self._browser.browser_session_id,
             browser_address=self._browser.browser_address,
+            request_options=RequestOptions(additional_headers={"X-User-Agent": "skyvern-sdk"}),
         )
+        LOG.info("AI workflow is running, this may take a while", run_id=workflow_run.run_id)
 
         workflow_run = await self._wait_for_run_completion(workflow_run.run_id, timeout)
         return WorkflowRunResponse.model_validate(workflow_run.model_dump())